TensorflowモデルをTFLiteにconvertする方法とconvertツールの比較

カスタムのTensorflowのモデルをTFLiteにconvertしようとしてすごく辛かったのではまりどころを記録していく。

サンプルにあるモデルをtfliteにconvertするのはそんなに難しくないんだが、ちょっと自分で手を加えたモデルをconvertしようとしたらTensorFlow初心者の私にはものすごく大変だった。

今回使うのはtensorflow/modelsのslimに入ってるモデル。 mobilenet、inception、resnet、vggなどのアーキテクチャが準備されている。 Kerasモデルのconvertはやったことがないので今回の話にはでてこない。

※最新のTensorflowのversionは2.2だが、tensorflow/modelsに準備されているモデルはTensorflow2には対応していない
※今回書いているPythonコード内のTensorfFlowはすべて2.2、Pythonは3.7.6
※コマンドラインツールやtocoは明記して別のバージョンを使っている

TFLiteのモデルにconvertするツールの比較
Checkpointからのconvertする
さいごに
おまけ
- トラブルシューティング
  - 解決したやつ
  - 解決できなかったやつ
- 超便利ツール

TFLiteのモデルにconvertするツールの比較

convertに使うツールいくつか候補がある。
推奨されてるかどうかではなくて自分が使いやすいと思った順にかいていく。多分推奨と逆の順番だけど。

toco

tocoは後述するtflite_convertやPython APIの元になっているプロジェクトで、元になってるだけあってオプションが一番多い。
自由自在にconvertしたいならこれが一番使いやすかった。

bazelで自分でbuildして使う。これを採用してconvertする方法を後述するので詳細な使い方はそちらで。

tflite_convert コマンド

tocoをtensorflowがwrapしたコマンドラインツールがtflite_convert。 tocoとそれほど変わらないが、tensorflowによってオプション名が変わっていたりする気がする。

公式は細かいこと設定したいならtflite_convert使ってねと言っているが、tf2系を入れていると、オプションが少なくて使えない気がする。このコマンドを使うなら、tf1系の環境を用意したほうが良い。

tflite_convertの元コードを見ると1系のオプションと2系のオプションの違いがわかる。

PythonAPI

ドキュメントに推奨と書いてあるのがPython API

v2のtf.lite.TFLiteConverter

推奨してる方法。saved_modelには対応してるけどfrozen GraphDefs形式の変換には対応していない。

converter = tf.lite.from_saved_model(saved_model_path)
converter.convert()

でsaved_modelのconvertをできるのでめっちゃ簡単(そうにみえる)。

でもtf.lite.TFLiteConverter.from_saved_modelのヘルプを見るとわかるけどオプション(引数)が以下の3つしかない。

saved_model_dir
signature_keys
tags

特に

inputのレイヤー・outputのレイヤーが指定
Quantization-aware training(inference_type、post_training_quantizeなど)のオプション

がないので推奨している割に一番使えない方法だと思う・・・。

v1のTFLiteConverter(tf.compat.v1.lite.TFLiteConverter)

Python APIはもとからレイヤーの指定をするオプションがないのかといったらそんなことはなくてv1のAPIのドキュメントにも書いてあるように、v1APIはv2APIに追加で以下の引数を渡すことができる。

input_arrays
input_shapes
output_arrays

こんなかんじに使える

converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(
    saved_model_path,
    input_shapes={"input_layer_name": [batch_size, width, height, channel]}
)
converter.convert()

ツールの紹介は以上にして、convertの仕方について書いていく。

Checkpointからのconvertする

Checkpointからconvertするには以下のような手順を踏む。 f:id:cocodrips:20200411182334p:plain:w400

GraphDefファイルを準備
モデルをFrozenする
モデルをTFLiteにconvert

今回はslim/netsの中にあるmobilenet_v2のmobilenet_v2_1.0_224_ckptをDLして、tocoを使ってconvertしていく。

slimの実装じゃなくても、GraphDefファイルがあればこの方法で問題ないと思う。

上記データにはfrozenGraphDef(frozen.pb)もtfliteモデルもはいってるけど、カスタムモデルを想定してこれを作るところからやっていく。

必要なリポジトリとライブラリのダウンロード

tensorflow/tensorflow

tocoはコマンドとして提供されていないのでtensorflowをclone。
r1.14のブランチにcheckoutする。
それ以前でも良いかもしれないけれど、1.15以降はtf.contribの実装がないのでslimのexporter(GraphDefを作るやつ)が動かない。

tensorflow/models

slimを使ってるモデルを想定してるのでもう入ってるとは思うけど、tensorflow/modelsをcloneする

bazel

tocoを使うのにgoogle製のビルドツールbazelを入れる必要がある。 brewとかでは入れないでバイナリから0.26をインストールする。最新のtensorflowはbazel2.0にupgradeされているっぽいけど、上記のtocoを使うには0.26が必要。

macでのinstall方法

GraphDefファイルを作る

slimのモデルはGraphDefファイルを作るexporterが実装されてるので、それを使ってexportする。

cd <tensorflow/models>
python research/slim/export_inference_graph.py \
  --alsologtostderr \
  --model_name=mobilenet_v2 \
  --dataset_name=imagenet \
  --batch_size=1 \
  --image_size=224 \
  --output_file=<model checkpoint dir>/mobilenet_v2_1.0_graph.pb

--model_nameはちがうモデルを指定したならここで対応を確認: nets_factory.py

FrozenGraphDefファイルを作る

せっかくカスタムモデルであることを想定しているので、最終層以外にいくつかのレイヤーを取り出すモデルにしてみる。

ラベルを推論したMobilenetV2/Predictions/Softmax
ラベル推論前の特徴量のMobilenetV2/Logits/AvgPool2D

の２つを取り出せるようにfrozenする。取り出したいレイヤーの名前を確認したいときはGraphDefファイルをnetronで開いてみると良い。

cd <tensorflow/tensorflow> 
bazel run tensorflow/python/tools:freeze_graph -- \
  --input_graph=<model checkpoint dir>/mobilenet_v2_1.0_graph.pb \
  --input_checkpoint=<model checkpoint dir>/mobilenet_v2_1.0_224.ckpt \
  --input_binary=true \
  --output_graph=<model checkpoint dir>/mobilenet_v2_1.0_224_frozen.pb \
  --output_node_names=MobilenetV2/Logits/AvgPool2D,MobilenetV2/Predictions/Softmax

tensorflow/python/tools:freeze_graphのビルドは長くて初回は手元のCPUが12コアあるPCで3時間くらいかかった。

FrozenGraphDefファイルをTF Liteモデルにコンバート

ここまでできてればあとは簡単

cd <tensorflow/tensorflow>
bazel run -c opt tensorflow/lite/toco:toco -- \
    --input_file=<model checkpoint dir>/mobilenet_v2_1.0_224_frozen.pb \
    --input_arrays=input \
    --output_file=<model checkpoint dir>/mobilenet_v2_1.0_224_float32.tflite \
    --output_arrays="MobilenetV2/Logits/AvgPool2D,MobilenetV2/Predictions/Softmax" \
    --inference_type=FLOAT \
    --allow_nonascii_arrays

これで完成！

でもこれだとconvertしたもののweightはfloat32でモデルは軽量化されていない。
TF Liteではfloat16とuint8にもconvertできるのでやってみよう。

uint8(モデルサイズ1/4)にquantizeするには上記コマンドに以下を加える。

--post_training_quantize

float16 (モデルサイズが1/2)にするには上記コマンドに以下を加える。

--post_training_quantize --quantize_to_float16

詳細はPost-training quantization のドキュメントへ。

追記(2020/04/27): float16にtfliteのinferencerが対応してるのは1.15以降っぽい: 該当コード

TF Liteモデルを使って推論

import matplotlib.pyplot as plt
from PIL import Image
import numpy as np

# image: https://www.flickr.com/photos/kneva/560380352, crop 224, jpg
image = np.array(Image.open('images/test.jpg'))
print(image.shape)

plt.imshow(image)
plt.show()

(224, 224, 3)

import pathlib 
import tensorflow as tf
print("tf version:", tf.__version__)

model_path = pathlib.Path(os.environ.get("CHECKPOINT_DIR")) / "mobilenet_v2_1.0_224_float32.tflite"
interpreter = tf.lite.Interpreter(model_path=str(model_path))
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# print(input_details)
# print(output_details)
    
interpreter.set_tensor(input_details[0]['index'], np.expand_dims(image, 0).astype(np.float32))
interpreter.invoke()

for output in output_details:
    print("===")
    print(output['name'])
    tensor = interpreter.get_tensor(output['index'])
    print(tensor)

tf version: 2.2.0-dev20200323
===
MobilenetV2/Logits/AvgPool2D
[[[[0.27684855 0.5009162  0.00454296 ... 0.         0.02394939
    0.6504333 ]]]]
===
MobilenetV2/Predictions/Softmax
[[3.4995192e-05 2.3994156e-05 1.5229588e-04 ... 1.8983837e-05
  8.8593506e-06 8.0088567e-04]]

こんなかんじで特徴量とラベルの推論値をTFLiteモデルから取得できた！

　　／⌒ヽ／⌒ヽ／⌒ヽ
⊂ヽ(　^ω^)つ^ω^)つ^ω^)つ
　＼　　／　　／　　／
　　( ＿フ( ＿フ( ＿_フ
　　(／　　(／　　(／

さいごに

TFLiteまわりは2系でsupportがまだまだな感じなので、今のところは1系使うのが良さそう。 TensorFlowさん、2系ゴリ押しで2系はドキュメントがあるんだけど1系の情報探そうとするとすごい大変だった・・・。

コンピュータビジョン専門でもないし、普段深層学習モデルをつくる仕事をしているわけでもないので、おかしな部分もあるかもしれないけれど、バージョン周りとか、ツールによるできることの違いとか、知見が色々あったのでせっかくなのでまとめた。

おまけ

トラブルシューティング

自分がハマったエラーメモしておく。

解決したやつ

tfliteにconvertてきたけどグラフが読み込めない

input->dims->size != 4 (0 != 4)Node number 0 (CONV_2D) failed to prepare.

今回の場合exportする際に--batch_size=1 を指定しそびれたせいでこのエラーがでた。 GraphDefファイルを作るときにbatch_sizeやinput_tensorを固定しないとこんな感じのエラーがでる。
TFLiteはinputのTensorをshapeを固定する必要がある。-1のままでは動かない。

frozenかexportしようとしたとき

ERROR: Linking of rule ‘//tensorflow/python:gen_sparse_ops_py_wrappers_cc’ failed (Exit 1)

bazelのversionがあってないとでた。 bazel含め環境が正しいか確認するにはtensorflowのディレクトリで./configureをちゃんとたたいて、自分の環境が間違ってないことを確認しよう。

解決できなかったやつ

PythonAPI経由でconvertしようとしたとき

Check failed: data_type == DT_FLOAT Fatal Python error: Aborted

で止まってconvertできなかったけど情報量がなさすぎてわからなかった

似てるissue:

超便利ツール

これ本当にfrozen成功してるのかな？
なんかconvertできたのに動かないなーー？
っていうのの繰り返しで、そんなときにめちゃくちゃデバッグに役立ったのがnetronというグラフの可視化ツール。

netron
↑さっき作ったtfliteファイルもこんな感じにひらける

Netron supports ONNX (.onnx, .pb, .pbtxt), Keras (.h5, .keras), Core ML (.mlmodel), Caffe (.caffemodel, .prototxt), Caffe2 (predict_net.pb, predict_net.pbtxt), Darknet (.cfg), MXNet (.model, -symbol.json), ncnn (.param) and TensorFlow Lite (.tflite).

とドキュメントにあるように、とにかくいろんな形式のモデルの可視化に対応している。
githubのリポジトリに9kくらいスターついてるし常識ツールなのかもしれないけど、これを知らないがために多大な時間を失った。
今までtensorflowのsession作ったりしながらtensorboardにグラフを出力したりして確認してたけど、そんな事する必要はなかったんだ・・・！