iMind Developers Blog

iMind開発者ブログ

TensorFlow Liteでモデルのサイズを小さくする(v1.14.0)

概要

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を指定して複数画像を同時実行

上述の例では1画像ずつ処理するフローになっている。複数レコードを同時に処理したい場合は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, 記事投稿