TensorFlow Liteを使ってモバイル向けに最適化された機械学習モデルを動かしてみる

2017
12.15

TensorFlow Liteを使ってモバイル向けに最適化された機械学習モデルを動かしてみる

こんにちは。GClueでアルバイトをしている木村です。普段は学生をしながらTensorFlow User Group Utsunomiyaという勉強会をやっています。ディープラーニングの分野では、実用アプリケーションにおける機械学習システム、組込み・モバイル向けのニューラルネットの高速化・軽量化技術に興味があります。

本記事は、TensorFlow Advent Calendar 2017 15日目の記事です。今回は、TensorFlow Liteの使い方について紹介します。

今年のGoogle I/Oで予告されていたTensorFlow Liteが、11月の中旬にデベロッパープレビューとして公開されました。TensorFlow Liteは既存のTensorFlowの拡張としてモバイル/組込み向けに最適化されたソリューションとされており、学習済みモデルのコンバータや推論用インタプリタ、そのためのC++/Java APIをカバーしています。TensorFlow Liteに関する情報は公式ドキュメントに“Introduction to TensorFlow Lite”として公開されていますが未だにデベロッパープレビュー版なので軽い紹介にとどまっています。詳細な情報を知るには、Githubのレポジトリのドキュメントを読む必要があります。

モバイル/組込みにおける機械学習システム

モバイル/組込みシステムにおいて機械学習システムを取り入れる場合、汎用コンピュータに比べて次ののような制約が想定されます。

計算性能（反応速度や計算コスト）
モデルのサイズ
消費電力

もちろんクラウド上の高性能コンピュータを使い推論を行うことも可能ですが、リアルタイム性が重要視されるシステムやプライバシーの観点からオンラインを避けるシステムでは、推論をローカルで行う必要があります。実際に、NVIDIAによる自動運転技術では30fps（画像1枚あたり0.033秒程度）という間隔で画像認識を処理することが想定されています。

以上のような問題のなか、現在では量子化ニューラルネット（Quantized Neural Networks）や蒸留（Distillation）といった理論的なアプローチが考えられていますし、NVIDIA JetsonやMovidius Neural Compute Stickなどの組込み向けに省電力・高速化を目的とするハードウェアからのアプローチも既に登場しています。

クロスプラットフォーム・省サイズ・高速を掲げるTensorFlow Liteはこのような背景のもと、ソフトとハードの進展を取り入れた形で登場したソリューションといえます。

TensorFlow Liteの使い方

公式ドキュメントによればTensorFlow Liteは、On-deviceでの機械学習の推論を低レイテンシと省バイナリサイズで行うこと可能とします。そのための技術としては、カーネルの最適化や量子化、FlatBuffersといった技術が使われています。

TensorFlow Liteの使い方は以下の図に示すように、1. 通常通りにモデルを学習させ、2. 学習済みモデルをTensorFlow Lite Converterにより変換し、3. 変換したモデルをモバイルプロジェクトに組み込む、という流れになります。

TensorFlow Lite Architecture (引用 https://www.tensorflow.org/mobile/tflite/ )

TensorFlowによる学習フェーズ
- 学習済みモデルの.pbファイルと.ckptファイルを作成します
TensorFlow Lite Converterによるモデルの変換
- TensorFlow Lite向けにモデルを変換し、.tfliteファイルを作成します
モバイルアプリ上でのモデルの推論
- .tfliteファイルをプロジェクトに組込み、C++/Java APIからInterpreterにより推論を実行します

TensorFlow Lite Converterの使い方

TensorFlow Lite Converter(toco)を使ってどの程度モデルが圧縮されるか確認してみましょう。

公式で用意されているMobileNetを使います。MobileNetは文字通りモバイル端末上で効率的に推論を行うことを目的とした畳み込みニューラルネット（CNN）で、パラメータ数を減らす工夫がなされています。

それではMobileNetの学習済みモデルを.lite形式のファイルに変換してみましょう。その前にtensorflowのソースコードからtocoというツールをビルドする必要があります。

$ git clone https://github.com/tensorflow/tensorflow.git
$ cd tensorflow
$ bazel build tensorflow/contrib/lite/toco:toco

TensorFlowのレポジトリ直下で以下のコマンドを実行します。これが、フリーズされた学習済みのモデルとなります。

$ wget https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_1.0_224_frozen.tgz
$ tar -xvf  mobilenet_v1_1.0_224_frozen.tgz

tensorflow/outputディレクトリにmobilenet_v1_1.0_224.liteファイルを作成します。以下の例では、FLOAT（浮動小数点数）を指定して変換しています。

$ mkdir output
$ bazel-bin/tensorflow/contrib/lite/toco/toco \
  --input_file=$(pwd)/mobilenet_v1_1.0_224/frozen_graph.pb \
  --input_format=TENSORFLOW_GRAPHDEF \
  --output_format=TFLITE \
  --output_file=$(pwd)/output/mobilenet_v1_1.0_224.lite \
  --inference_type=FLOAT \
  --input_type=FLOAT \
  --input_arrays=input \
  --output_arrays=MobilenetV1/Predictions/Reshape_1 \
  --input_shapes=1,224,224,3

次に、QUANTIZED_UINT8を指定して量子化してみましょう。公式ドキュメントによれば、ReLu6を使う場合、フラグ--default_ranges_minと--default_ranges_maxを指定して「ダミー量子化」する必要があります。

$ bazel-bin/tensorflow/contrib/lite/toco/toco \
  --input_file=$(pwd)/mobilenet_v1_1.0_224/frozen_graph.pb \
  --input_format=TENSORFLOW_GRAPHDEF \
  --output_format=TFLITE \
  --output_file=$(pwd)/output/mobilenet_v1_1.0_224_quantized.lite \
  --inference_type=QUANTIZED_UINT8 \
  --input_array=input \
  --output_array=MobilenetV1/Predictions/Reshape_1 \
  --input_shape=1,224,224,3 \
  --default_ranges_min=0 \
  --default_ranges_max=6 \
  --mean_value=127.5 \
  --std_value=127.5

実際に生成されたファイルのサイズを比較してみます。学習させた素の学習モデル(frozen_graph.pb)に比べて、FLOATを使って変換したモデル(mobilenet_v1_1.0_224.lite)ではサイズは変わりませんでしたが、QUANTIZED_UINT8を指定して量子化したもの(mobilenet_v1_1.0_224_quantized.lite)はサイズがおよそ4分の1となりました。したがって、量子化により学習済みモデルがスリムになってことが分かりました。ちなみに、Appleが公開しているCoreMLによるMobileNetのモデルサイズは17.1MBとなっています。

	サイズ(MB)
`frozen_graph.pb`	16
FLOAT( `mobilenet_v1_1.0_224.lite)`	16
QUANTIZED_UINT8 ( `mobilenet_v1_1.0_224_quantized.lite)`	4.1

AndroidでMobileNetを動かす

TensorFlowのレポジトリ直下のtensorflow/contrib/lite/java/demoをAndroid Studioで開きます。なお、Android SDKのバージョンは26以上（Android Oreo）、Android NDKは14以上が必要となります。

公式で公開されているモデルファイルをダウンロードし、mobilenet_quant_v1_224.tfliteをtensorflow/contrib/lite/java/demo/app/src/main/assets/に配置してください。

$ cd tensorflow/contrib/lite/java/demo/app/src/main/assets/
$ wget https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip
$  unzip mobilenet_v1_224_android_quant_2017_11_08.zip
$ rm mobilenet_v1_224_android_quant_2017_11_08.zip

ビルドする準備が整ったので、MobileNetを動かしてみましょう。今回使用した端末は、LGE Nexus 5X (Android 8.1.0, API 27)です。

以下は実行例です。ハサミとパソコンをキャプチャしてみました。ハサミの場合、spatula（へら）が最上位で出ていますがletter openerも候補に挙がっています。パソコンのキーボードの場合、keyboardが最上位に出ていることが確認できます。

まとめ

以上、TensorFlow Liteの基本的な使い方を公式ドキュメントに沿いながら補足を入れて紹介してみました。今回TensorFlow Liteを触ってみて、既存の学習済みの機械学習モデルの資産を無駄にせずモバイル・組込みに最適化した形で再利用できることが、TensorFlow Liteの最大のメリットだと思いました。今後、TensorFlow MobileからTensorFlow Liteに切り替わっていくことが公式でアナウンスされているので、みなさんも今から使ってみてはいかがでしょうか。