TensorFlow Lite を使ってみた。

概要

TensorFlow Liteを使って、TensorFlowで書かれたモデルを変換し、推論してみました。
その際に詰まったところや、計測時に出た疑問点を書いておきます。

他の方も同じようなことをされていましたが、TensorFlow や Python のバージョンの違いのせいか、うまくいかなかったので、それも比較しながら書いていきます。(以下、旧バージョンと書きます)
比較させていただいたサイト: https://qiita.com/yohachi/items/434f0da356161e82c242#%E5%8F%82%E8%80%83

環境

OS: ubuntu 16.04 LTS, Intel® Core™ i5-7200U CPU @ 2.50GHz × 4
Python: 3.5.2
TensorFlow: 1.12.0 (計測環境),　1.13.1（その他の準備環境など）
Edge TPU: Google Coral Edge TPU USB Accelerator [google-002]

TensorFlow Lite　とは

エッジ端末向けの TensorFlow の改良版。通常の TensorFlow で書かれたモデルの重みを8bitに変換して推論する。
学習自体は TensorFlow で行われる。その際のチェックポイントファイル(.pb、.h5 などの形式）を TensorFlow Lite 用のファイル(.tflite)に変換してロードすることで、TensorFlow Lite で推論可能。

.tfliteファイルへの変換

旧バージョンでは、どうやら次のようにかくと、.h5ファイル（keras や tf.keras を用いて書かれた重み付きモデルの形式。重みファイルのみを保存している場合はエラーとなることに注意）を .tflite 形式に変換できるようだ。（ターミナル上でコマンドで変換することも可能）

convert.py

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_keras_model_file("keras_model.h5")
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)

しかし、私の環境で実行したところ、以下のようなエラーが発生した。なお、旧バージョンでは TensorFlow の Tensorflow Lite のためのクラスが tensorflow クラス直下に存在するが（tf.lite.~)、tf.contrib.lite.~ に移動しているので、そこは修正しています。

$ python3.5 convert.py 
Traceback (most recent call last):
  File "convert.py", line 3, in <module>
    converter = tf.contrib.lite.TFLiteConverter.from_keras_model_file("model.h5")
  File "/home/mori/.local/lib/python3.5/site-packages/tensorflow/contrib/lite/python/lite.py", line 368, in from_keras_model_file
    keras_model = _keras.models.load_model(model_file)
  File "/home/mori/.local/lib/python3.5/site-packages/tensorflow/python/keras/engine/saving.py", line 230, in load_model
    model = model_from_config(model_config, custom_objects=custom_objects)

raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
ValueError: Unknown layer: BatchNormalizationV1

原因を調べてみたところ、tf.keras expoter でのエラーのせいでこのようなエラーが出る様子。変換ではなく、重み付きモデルファイルを読み込む段階でエラーが出るため、tf.keras.models.load_model(modelfile)を実行してももちろん同様のエラーが発生する。参考: https://github.com/tensorflow/tfjs/issues/1334
なお、tf.keras ではなく、ノーマルの keras を使ったり、 TensorFlow 2.0.x を使ったりすると上記の方法で大丈夫なようだ。（未検証）

その解決策

.h5 形式のファイルを .tflite 形式に変換ができないので、 .pb 形式（グラフ）で保存してやり、それを直接 .tflite 形式に変換する。

再度学習させて .pb 形式で保存するのが面倒だったので、下記の要領で .h5 形式を .tflite 形式に変換させることにした。よって、

1. .h5 ファイルを出力（学習の段階でチェックポイントとして）
2. .h5 ファイルを読み込み、モデルをグラフとして保存(今回は、重みファイルのみを保存していたので、モデルに一旦ロードする必要がある。)
3. 変換されたグラフを .tflite 形式に変換
4. 変換された .tflite 形式のファイルを読み込み

2から順番に、

convert_h5_to_pb.py　のmain関数

if __name__ == "__main__":
    load_path = "./model_weights_999.h5" # Pre-trained
    model = getModel(640, 480, 3, 3, load_path)
    test_gen = Generator(batch_size=64, data_num=1000, x_memmap_path="./test_X.npy", y_memmap_path="./test_y.npy")

    #pred = model.predict_generator(test_gen)
    #expect = test_gen.memmap_y
    #print_accuracy(pred, expect)

    #tf.keras.experimental.export_saved_model(model, './saved_models') # Able to apply with the latest version of Tensorflow                       
    tf.contrib.saved_model.save_keras_model(model, "./saved_models")

convert_pb_to_tflite.py

import tensorflow as tf
import sys

try:
    tf_conv = tf.contrib.lite.TFLiteConverter.from_saved_model("saved_models/1582435994/")

    #Fail                                                                                                                                 
    #tf_conv = tf.contrib.lite.TFLiteConverter.from_keras_model_file('model.h5')                                                          

    #Fail                                                                                                                                 
    #new_model = tf.keras.models.load_model("./model.h5")                                                                                 

    print("\n\nSuccessfully loaded\n\n")
    lite_model = tf_conv.convert()

except Exception as e:
    print("\n\nFailed to load\n\n", e)
    sys.exit()

with open('digit_model.tflite', 'wb') as w:
    w.write(lite_model)

convert_pb_to_tflite.py の5行目において、変なディレクトリを指示しているのはモデルを保存する際の仕様がもともとそうなっているからです。

$ tree saved_models	
saved_models/
└── 1582256034
    ├── assets
    │   └── saved_model.json
    ├── saved_model.pb
    └── variables
        ├── checkpoint
        ├── variables.data-00000-of-00001
        └── variables.index

3 directories, 5 files

これで、無事 .tflite 形式にモデルを変換することができた！

計測

TenforFlow Lite にせっかく変換したのだから、比較してみます。比較対象は次の２つ。

1. 推論の実行速度　(10回実行した平均値、推論のみの時間を計測)
2. モデルのサイズ
3. 推論の精度

比較するモデル、環境は次の3つ。
A. 計測環境で通常の .h5 形式の重み付きモデルを使って推論（32 bits）
B. 計測環境で、.tflite 形式の重み付きモデルを使って推論 (8 bits)
C 及び C'. Coral (Google Coral Edge TPU USB Accelerator [google-002])で、 .tflite 形式の重み付きモデルを使って推論 (8 bits)

なお、もともと tf.keras を用いて作成したネットワークのモデルは、Dronet　(http://rpg.ifi.uzh.ch/dronet.html)
を修正したものです。画像を入力とし、次に動くべき方向（Left, Forward, Right）の one-hotベクトルを出力します。

テストデータとして、1000枚の画像、1枚の画像を入力として与えた際の両方で比較しました。

テストデータ=1000

	実行速度 (ms)	モデルのサイズ (KB)	出力の精度（％）
A	40,525.3	1,793	82.2
B	59,289.5	1,704	82.2
C	41,329.2	1,704	82.2

テストデータ=1

	実行速度 (ms)	モデルのサイズ (KB)	出力の精度（％）
B	67.3	1,704	100.0
C	19,650.2	1,704	100.0
C'	57.5	1,704	100.0

.tflite 形式にしても、そんなにモデルの大きさが変わらなかったのは、もともとのモデルが小さいからと思います。
っていうか、テストデータが画像１枚のときの C の計測間違ってないか？と思われたかもしれませんが、これはコードの書き方に問題でした。テストデータの枚数を引数として与えていたので、同じ関数を使いまわしていたのですが、その際のループの書き方が Coral 的には良くないみたいです。下記の修正後の実行を C' としています。

run.py（Cで計測したコードの一部）

    start = time.perf_counter()
    for im in images:
        #im = np.reshape(im, [1, 480, 640, 3])
        coral.set_input(interpreter, im)
        interpreter.invoke()
    inference_time = time.perf_counter() - start
    print('%.1fms' % (inference_time * 1000))
    return inference_time

run.py（C'で計測したコードの一部）

    start = time.perf_counter()
    #im = np.reshape(im, [1, 480, 640, 3])
    coral.set_input(interpreter, im)
    interpreter.invoke()
    inference_time = time.perf_counter() - start
    print('%.1fms' % (inference_time * 1000))
    return inference_time

ここ以外の部分は修正していません。なお、

run.py（C'で計測したコードを修正）

    start = time.perf_counter()
    for _ in range(1):
        coral.set_input(interpreter, images)
        interpreter.invoke()
        inference_time = time.perf_counter() - start
    print('%.1fms' % (inference_time * 1000))
    return inference_time

のように、ループを定数で固定して回してやると、同等の実行速度が得られました。Coral は多次元配列のループに弱いのかと思いきや、テストデータが1000枚のときはこれでうまくいっているので謎です。しっかり実行速度もだいたい1000倍になってますし。、公式サイトを調べようとしましたがここで力尽きました。

ちなみに、精度の計算方法ですが、全てのテストデータの正答率を最終的な精度としています.

Calc    : 
 [[230.   0.  35.]
 [ 56.   0.  55.]
 [ 32.   0. 592.]]
Labels  : 
 [[265.   0.   0.]
 [  0. 111.   0.]
 [  0.   0. 624.]]
Average_accuracy : 
 [0.86792453 0.         0.94871795]

Accuracy:  0.822

こんな感じで3次元版の混合行列を作成してみましたが、圧倒的に2つ目のラベルの正答率が悪いですね。テストデータの配分から見て頂いてわかるように、学習用データの数も少なかったからでしょうか。

この出力が、3つのモデル、環境全てで一致したので、テーブルの中で精度がすべて同じになっていますが、これはおそらく全てのテストデータ内の画像において同じ出力を返していると考えて大丈夫そうです。

考察

わりと単純なモデルだったせいか、量子化しても（32bits → 8bits）モデルのサイズがほとんどかわりませんでした。これだけでは量子化のメリットがそこまでないですね。
ですが、精度はなんと全く変わらないという結果でした。これくらい層の少ないモデルであれば量子化は問題ないと考えていいのでしょうか。

Coral と比較してみた結果、なんと通常のモデルでの推論とほぼ変わらんやんけと文句を頂戴しそうですが、テストデータが1000枚の際、推論をミニバッチ(batch_size=64)処理で行っているため、そこで高速化されているのかなと思いました。同じように1000回ループを回したら Coral 上で実行したほうがはやそうですね。それで比較しろよと言われそうですが、わざわざミニバッチを使わない利点もないうえに、実際に推論する際は１枚くらいしか入力がないケースしかなさそうなのでとりあえずこれで満足しています。（修正して追記をあげる可能性有り）

今後の目標

今回の比較は、ニューラルネットワークのパラメータを量子化してみようということで行われました。8ビットなら精度が変わらないということで、最終的に1ビットにできたらなぁとか画策していたりします(BNN)。今回は自分で量子化する手間を省いて TensorFlow Lite を使いましたが、もっとパラメータを小さいビット数に量子化しようと思うと、自分で一からコードを書くしかないのかなぁと思っています。
そういったライブラリや技術などの知識をお持ちの方は、ぜひ教えてください！飛び跳ねて喜びます。

あとは、学習用とテスト用データをちゃんと揃えることですね。今のままだと2つ目のラベルの正答率が低すぎます。
それに加えて、教師データを用意する際に、実は手動で調整されたスクリプトを使用しています。なので実際に期待される出力とは異なっている可能性があるので、これもちゃんとしないとです。

今回はネットワークの出力が方向の one-hotベクトルでしたが、これに物体との衝突確率や、停止指示なども出力できるようにしたいですね。というか参考にさせていただいたネットワークにはついていました。（笑）

雑談

これが初めて書かせていただいた記事になります。見づらい、こういう構成のほうがいいよ等意見や感想がございましたら、教えていただけると幸いです。今後もネットワークの量子化や、その実装について記事を書いていけたらなぁと思ってますので、よろしくお願いします。