Tensorflow Lite のモデルをAndroid端末で動かす

彌冨研 Advent Calendar 2019 22日目の記事です。
当日の夜19時から始めてその日のうちにTFLiteのモデルをスマホにデプロイするまでの話。

デプロイは当日中にできたが、記事は当日中に書ききれなかった。。。ごめんなさい。

これまでのあらすじ

“僕はロボットのための認識器をつくりたい” をよんでね

エッジデバイスでNNモデルを推論させよう

エッジデバイスで推論させるためにNNを軽量にしたり推論時の速度を高速にするテクニックが色々とあります。
基本的に、各アプローチは以下の3わけられるはずです。

  • Convolutionの演算回数を減らす
    • Group / Depthwise Convolution
    • ShuffleNet
  • パラメータの量子化
  • 推論向けのグラフ最適化

“Convolutionの演算回数を減らす” と “パラメータの量子化” は実際に推論するデバイスによって効果が変わってきます。
手法によっては、FLOPSは確かに減っているんだけども並列性が犠牲になっていてGPU推論に向かなかったり、
デバイスによっては、16bit floatや8bitに最適化されきっていなかったりします。

しかしながら、速度だけ見ていればよいかというとそういうわけでなく、
スマホのようにバッテリーを持つデバイスであれば消費電力も考慮するべきです。
そういうときには、量子化のようにパラメータ数も演算時のプロセッサへの負荷も大きく減らせる方法が有力になるでしょう。

Tensorflowモデルをスマホで動かすまで

さてさて、本題です。
エッジ推論用のフレームワークはTVM (ワシントン大), SNPE (クアルコム), ONNX Runtime (マイクロソフト) など色々ありますが、今回はGoogle製のTensorflow Liteを使うことにしました。

  • Tensorflowの学習済みモデルの多くが使える
  • スマホ向けのexample projectが充実している
  • 量子化機能がついている

下準備

Ubuntu 18.04で作業してます。

python3 --version #3.6.8
pip3 install tensorflow==1.14.0
wget https://dl.google.com/coral/python/tflite_runtime-1.14.0-cp36-cp36m-linux_x86_64.whl
pip3 install tflite_runtime-1.14.0-cp36-cp36m-linux_x86_64.whl 
wget -O protobuf.zip https://github.com/google/protobuf/releases/download/v3.0.0/protoc-3.0.0-linux-x86_64.zip
unzip protobuf.zip

TensorflowモデルをTFLiteモデルに変換する

今回はMobileNetv1 SSDモデルを量子化してスマホで動かします。

本当はMobileNetv2でやりたかったのですが、未だに動かせないというIssueがOpenしている。。

まずはTensorflow Object Detection APIのインストールしておく

git clone https://github.com/tensorflow/models.git
cd models/research
python3 setup.py build
python3 setup.py install
../../bin/protoc object_detection/protos/*.proto --python_out=.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

モデルを落としてきて、TFLite形式のモデルに変換したあとにTensorflow Lite Optimize Converter (TOCO) で量子化などの推論向け最適化を実施する。

wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2018_01_28.tar.gz
tar -zxvf ssd_mobilenet_v1_coco_2018_01_28.tar.gz 
python3 object_detection/export_tflite_ssd_graph.py \
    --pipeline_config_path=ssd_mobilenet_v1_coco_2018_01_28/pipeline.config \
    --trained_checkpoint_prefix=ssd_mobilenet_v1_coco_2018_01_28/model.ckpt \ 
    --output_directory=/tmp/tflite \ 
    --add_postprocessing_op=true
toco \ 
    --graph_def_file=/tmp/tflite/tflite_graph.pb \
    --output_file=/tmp/tflite/detect.tffile \
    --input_shape=1,300,300,3 \
    --input_arrays=normalized_input_image_tensor \
    --output_arrays='TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3' \
    --inference_type=QUANTIZED_UINT8 \
    --mean_values=128 --std_dev_values=128 \
    --change_concat_input_ranges=false --allow_custom_ops

スマホにのせる

ここはガッツリサボりました。 Tensorflow Exampleslite/examples/detection に物体検出をAndroidアプリとして走らせるプロジェクトがあるのでそのまま流用しちゃいました。
asset の中のモデルを入れ替えるだけです。

手持ちのPixel3で40msec.くらいで動いているのがわかります。
なにもチューニングしていないので、バウンディングボックスが二重に出ちゃったり暴れたりしてますね。。。
それでもTVやPerson、Mouseなどはちゃんと認識されていることがわかります。

オプションコマンド名が公式docと違うとか、Android Studioの設定で若干躓いたりとかしましたが、
19時に始めて23時くらいにはスマホ上で物体検出モデルが動いちゃんですから普通に便利です。
スマホアプリ部分をカスタマイズすれば、自分で学習した独自モデルなんかも動かせそうです。