概要
TensorFlow（Keras）で作ったモデルをTensorFlow Liteでコンバートする。
バージョン情報
- tensorflow==1.14.0
この資料でやること
TFLiteConverterには下記の4つのメソッドが用意されている。
- from_frozen_graph
- from_keras_model_file
- from_saved_model
- from_session
本稿ではこの中からfrom_keras_model_file, from_saved_modelの2つを利用する。
モデルの生成
適当なサイズのモデルを作る。
import tensorflow as tf from tensorflow.keras import layers from sklearn import model_selection # データのロード (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data() x_train = x_train.reshape([60000, 28, 28, 1]) x_test = x_test.reshape([10000, 28, 28, 1]) # testデータの1/10をvalidationデータに使う x_test, x_val, y_test, y_val = model_selection.train_test_split( x_test, y_test, test_size=0.1, stratify=y_test) # モデルの用意 model = tf.keras.models.Sequential([ layers.Convolution2D(64, (4, 4), activation='relu', input_shape=(28, 28, 1)), layers.MaxPooling2D((2, 2)), layers.Dropout(0.1), layers.Convolution2D(64, (4, 4), activation='relu'), layers.MaxPooling2D((2, 2)), layers.Dropout(0.3), layers.Flatten(), layers.Dense(256, activation='relu'), layers.Dropout(0.5), layers.Dense(64, activation='relu'), layers.BatchNormalization(), layers.Dense(10, activation='softmax') ]) model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy']) # tensoboard callback tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir="logs") # fit model.fit( x_train, y_train, epochs=200, validation_data=(x_val, y_val), callbacks=[tensorboard_callback])
余談。
今回の記事はTensorFlow 1.13.1でコードを書いていたけどcheckpointの指定とかをしていたら下記のエラーに遭遇する。1.14.0では直っている。
RuntimeError: Unable to create link (name already exists)
from_keras_model_file
from_keras_model_fileは名前の通りkerasでsaveしたモデルに対して実行できる。
まずは生成したモデルをsave。
model.save('digit_model.h5')
生成されたファイルは4.1MBだった。
これをtfliteでconvert。
tf_conv = tf.lite.TFLiteConverter.from_keras_model_file('digit_model.h5') lite_model = tf_conv.convert() with open('digit_model.tflite', 'wb') as w: w.write(lite_model)
保存されたファイルの容量は 4.1MB → 1.4MB と減少。
tfliteのモデルでpredictを実行
変換されたモデルを読み込んでaccuracyを比較してみる。
まずは変換前のモデルで確認。
del model model = tf.keras.models.load_model('digit_model.h5') y_pred1 = model.predict(x_test) y_pred1 = np.argmax(y_pred1, axis=1) import numpy as np sum(y_pred1 == y_test) / len(y_test) #=> 0.9933333333333333
99.3%だった。
続いてTFLiteConverterで変換したファイルに対して実行。
input_detailsのshapeは未指定の場合は(1, 28, 28, 1)になっているので、1つずつデータを渡してpredictしていく。
# 読込 interpreter = tf.lite.Interpreter(model_path="digit_model.tflite") interpreter.allocate_tensors() # 入出力のdetail確認 input_details = interpreter.get_input_details() #=> [{'name': 'conv2d_6_input', 'index': 7, 'shape': array([ 1, 28, 28, 1], dtype=int32), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}] output_details = interpreter.get_output_details() #=> [{'name': 'dense_11/Softmax', 'index': 16, 'shape': array([ 1, 10], dtype=int32), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}] # テストデータ（9000件）に対してpredict実行 from tqdm import tqdm def predict(i): input_data = x_test.astype(np.float32)[i:i+1] interpreter.set_tensor(input_details[0]['index'], input_data) interpreter.invoke() return interpreter.get_tensor(output_details[0]['index'])[0] y_pred2 = np.array([predict(i) for i in tqdm(range(len(y_test)))]) y_pred2 = np.argmax(output_data, axis=1) sum(y_pred2 == y_test) / len(y_test) #=> 0.9933333333333333
同じく99.3%。中身を見てみるとまったく同じ結果。
メモリ使用量の確認
使用しているメモリサイズも確認。
下記のような処理で使用メモリを取ってみる。psutil v5.6.3を利用。
def show_memory_info(): ''' メモリ使用量確認 ''' process = psutil.Process(os.getpid()) print(process.memory_info())
| rss | vms | shared | data | |
|---|---|---|---|---|
| 変換前 | 2,501,242,880 | 29,616,340,992 | 567,939,072 | 2,949,496,832 |
| 変換後 | 885,657,600 | 3,651,903,488 | 144,904,192 | 1,286,762,496 |
バッチサイズの差もあって物理メモリ使用量が1/3くらいまで抑えられている。その分、実行時間はかかるけど。
input_shapesを指定して複数画像を同時実行
上述の例では１画像ずつ処理するフローになっている。複数レコードを同時に処理したい場合はTFLiteConverterにinput_shapesを指定する。
下記は先頭に10を指定。
tf_conv = tf.lite.TFLiteConverter.from_keras_model_file(
'digit_model.h5',
input_shapes={'conv2d_input': (10, 28, 28, 1)})
lite_model = tf_conv.convert()
with open('digit_model.tflite', 'wb') as w:
w.write(lite_model)
10件まとめて実行。
def predict(i): input_data = x_test.astype(np.float32)[i:i+10] # 10件ずつ渡す interpreter.set_tensor(input_details[0]['index'], input_data) interpreter.invoke() return interpreter.get_tensor(output_details[0]['index'])
from_saved_model
from_saved_modelはsaverで保存したモデルに対して実行する。
Kerasのモデルのままでは使えないのでツールを使って変換。
# kerasのモデルをsaved_modelに変換 tf.keras.experimental.export_saved_model(model, 'digit_model')
こんな感じの見慣れた形式で出力される。
$ tree digit_model
digit_model
├── assets
│ └── saved_model.json
├── saved_model.pb
└── variables
├── checkpoint
├── variables.data-00000-of-00002
├── variables.data-00001-of-00002
└── variables.index
あとは読み込むだけ。
# from_saved_modelの実行 tf_conv = tf.lite.TFLiteConverter.from_saved_model( 'digit_model')
改定履歴
Author: Masato Watanabe, Date: 2019-08-11, 記事投稿