TensorFlow Lite / モバイルで実行するためのスタイル転送の最適化

以下の記事を参考に書いてます。

・Optimizing style transfer to run on mobile with TFLite

1. はじめに

「スタイル転送」は、「コンテンツ画像」（建物など）と「スタイル画像」（画家によるアートワークなど）の2つの画像から、それらをブレンドして、出力画像を生成する手法です。今回、モバイル向けに最適化された学習済みのスタイル転送モデルと、Android・iOSのサンプルアプリを共有できることを嬉しく思います。

この記事では、モバイル用にTensorFlow Liteモデルを最適化する過程と、TensorFlow Liteモデルをモバイルアプリで効率的に使用する方法について説明します。

2. バックグラウンド

「スタイル転送」は、「A Neural Algorithm of Artistic Style」で最初に公開されました。ただし、元の手法は計算コストが高く、ハイエンドのGPUでもスタイル転送に数秒かかりました。その後の研究によって、高速化する方法が示されました。

私たちは、いくつかのモデルアーキテクチャを評価した後、サンプルアプリ用に事前訓練された「Magenta」のスタイル転送モデルから始めることにしました。モデルは、「コンテンツ画像」と「スタイル画像」を入力として受け取り、フィードフォワードニューラルネットワークを使用して、様式化された「出力画像」を生成します。このモデルでは、Gatysの論文の手法と比較してスタイル転送がはるかに高速ですが、それでもかなり大きく（44 MB）、低速です（Pixel 4 CPUで2340 ms）。そのため、モバイルアプリでの使用に適するように、モデルを最適化する必要があります。この記事では、そうした経験と、仕事で活用できるリソースを紹介します。

3. モデルアーキテクチャの最適化

「Magenta」のスタイル転送モデルは、2つのサブネットワークで構成されています。

・Style prediction network：スタイル画像をスタイル埋め込みベクトルに変換します。
・Style transform network：スタイルの埋め込みベクトルをコンテンツ画像に適用して、様式化された画像を生成します。

「Style prediction network」は「InceptionV3」バックボーンだったため、モバイル用に最適化された「MobileNetV2」バックボーンに置き換えました。「Style transform network」は、いくつかの畳み込み層で構成されています。MobileNetの幅乗数のアイデアを適用し、すべての畳み込み層の出力チャネルの数を4倍に縮小しました。

次に、モデルの訓練方法を決定する必要がありました。ゼロからモバイルモデルを訓練するか、学習済みのMagentaのモデルから蒸留して、複数のオプションを試してみました。最初から他のパラメータを最適化しながら「MobileNetV2」の重みを修正すると、最良の結果が得られることがわかりました。

モデルを大幅に縮小および高速化しながら、同様のレベルのスタイルとコンテンツの損失を達成することができました。

4. 量子化

モデルアーキテクチャを決定したら、「TensorFlow Model Optimization Toolkit」を使用した量子化により、モデルをさらに縮小します。

これは、モデルのサイズを最大4倍に縮小し、わずかな品質のトレードオフでモデルの推論を高速化できるため、TensorFlowモデルのほとんどに適用できる重要な手法です。

TensorFlowが提供する利用可能な量子化オプションの中で、シンプルさとモデル品質の適切なバランスがあるため、訓練後の整数量子化を使用することにしました。TensorFlowモデルをTensorFlow Liteに変換するときに必要なのは、訓練データセットのごく一部です。

量子化後のモデルは、同じレベルのスタイルとコンテンツの損失を維持しながら、元のモデルよりも1桁以上小さく、高速です。

5. モバイルへの展開

Androidアプリを実装して、スタイル転送モデルの使用方法を示しました。アプリは「スタイル画像」と「コンテンツ画像」を受け取り、スタイルとコンテンツを組み合わせた画像を出力します。

携帯電話のカメラを使用して、Camera2 APIで「コンテンツ画像」をキャプチャし、有名な絵画のセットを「スタイル画像」として使用します。「コンテンツ画像」に「スタイル」を適用するには2つのステップがあります。まず、「Style prediction network」を使用して、配列を浮動小数点数としてスタイルを抽出します。次に、「Style transform network」を使用して、この「スタイル」を「コンテンツ画像」に適用します。

CPUとGPUの両方で最高のパフォーマンスを実現するために、チップごとに最適化されたTensorFlow Liteモデルの2つのセットを作成しました。 CPU推論にはint8量子化モデルを使用し、GPU推論にはfloat16量子化モデルを使用します。 GPUは一般にCPUよりも優れたパフォーマンスを実現しますが、現在サポートされているのは、int8量子化モデルよりも大きい浮動小数点モデルのみです。次に、int8モデルとfloat16モデルの動作を示します。

モバイルアプリでスタイル画像の固定セットのみをサポートすることを計画している場合は、「Style prediction network」の結果をキャッシュすることで、パフォーマンスが向上する可能性もあります。 これにより、ネットワーク全体のサイズの91％を占める「Style prediction network」を含める必要がないため、アプリが小さくなります。 これが、プロセスが1つだけではなく2つのモデルに分割される主な理由です。

サンプルはGitHubにあり、スタイルを適用するメインクラスはStyleTransferModelExecutorです。計算コストがかかるため、UIスレッドでスタイル転送を実行しないことが重要です。 代わりに、AndroidXおよびコルーチンのViewModelクラスを使用して専用のバックグラウンドスレッドで実行し、ビューを簡単に更新します。 また、GPUデリゲートを使用してモデルを実行する場合、TensorFlow Liteインタープリタの初期化、GPUデリゲートの初期化および推論は、すべて同じスレッドで実行する必要があります。

6. 本番でのスタイル転送

「Google Arts＆Culture」アプリは最近、TensorFlow Liteを使用してデバイス上でスタイル転送を実行する「Art Transfer」を追加しました。 使用されるモデルは上記のものと非常に似ていますが、速度とモデルサイズよりも品質を優先します。 プロダクションでスタイルの転送を確認したい場合は、ぜひお試しください。

7. あなたの番

自分のアプリにスタイル転送を追加したい場合は、モバイルサンプルをダウンロードすることから始めます。 TensorFlow Hubでは、float16（予測ネットワーク、変換ネットワーク）とint8量子化バージョン（予測ネットワーク、変換ネットワーク）の両方のモデルバージョンを使用できます。私たちはあなたが作成できるものを見るのを待つことができません。あなたの作品を私たちと共有することを忘れないでください。

8. リソース

端末で機械学習モデルを実行すると、ユーザーデータをプライベートに保ちながら、低レイテンシを提供できるという利点があります。

この投稿では、「TensorFlowモデル」を「TensorFlow Lite」に直接変換することが最初のステップにすぎない可能性があることを示しました。優れたパフォーマンスを達成するには、開発者は量子化を使用してモデルを最適化し、モデルの品質、モデルのサイズ、および推論時間の間の適切なトレードオフを見つける必要があります。

以下のリソースを使用してモデルを作成しました。端末上の機械学習のユースケースにも適用できる場合があります。

◎ Magenta model repository
「Magenta」はTensorFlowを利用したオープンソースプロジェクトです。
機械学習を使用して音楽やアートを作ります。
スタイル転送モデルを含め、「TensorFlow Lite」に変換できるモデルが多数あります。

◎ TensorFlow Model Optimization Toolkit
「TensorFlow Model Optimization Toolkit」は、量子化やプルーニングなど、モデルを最適化する複数の方法を提供します。

◎ TensorFlow Lite delegates
「TensorFlow Lite」は、GPUやDSPなど、デバイスで利用可能なさまざまなタイプのハードウェアアクセラレータを活用して、モデルの推論を高速化できます。