TensorFlow LiteをKerasモデルで試してみる

エッジAIに関する調査の続きということで、TensorFlow Liteを試してみました。

エッジデバイスは基本的に計算性能が貧弱なので、別の環境で作成した学習済みモデルを軽量化し、推論のみをエッジ側で行うというソリューションです。TensorFlowで作成したモデルをLite用に変換するという方法を取ります。

TensorFlowモデルを作るのはなかなかかったるい(TensorFlow 2.0だとそうでもないみたいですが)ので、いつもどおりKerasでモデルを作って、それをTensorFlow Lite用に変換していきます。

KerasでCNNモデルの作成

まず、データセットの準備。いつもどおりのMNISTです。

from keras.datasets import mnist
from keras.utils import to_categorical

(X_train, y_train), (X_test, y_test) = mnist.load_data()

# shape[0]はデータの個数、(28, 28)を(28, 28, 1)に
X_train = X_train.reshape(X_train.shape[0], 28, 28, 1)
X_test = X_test.reshape(X_test.shape[0], 28, 28, 1)

# データの正規化
X_train = X_train / 255
X_test = X_test / 255

y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)

次にKerasでCNNモデルを作ります。極めて単純なモデル。

from keras.models import Sequential
from keras.layers import Dense, Activation, Conv2D, MaxPooling2D, Flatten

model = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3), padding='same', input_shape=(28, 28, 1)))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(128))
model.add(Activation('relu'))
model.add(Dense(10))
model.add(Activation('softmax'))

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

学習します。今回は簡単にするためエポックは5。MacBook Proだと1分半くらいで学習が終わります。

# 学習
model.fit(X_train, y_train, batch_size=32, epochs=5, verbose=1)

# 評価
loss, accuracy = model.evaluate(X_test, y_test, verbose=0)
print('Accuracy', '{:.2f}'.format(accuracy))

で、ここで作成したKerasモデルをh5ファイルに保存します。

model.save('mnist.h5')

ファイルサイズを見てみると、9.2MBあります。

TensorFlow Liteモデルに変換

モデルの変換もPythonでできます。

import tensorflow as tf

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

ファイルサイズを比べてみると、だいたい3分の1になっていることが分かります。

TensorFlow Liteモデルで推論

次に推論処理を行います。これはRaspberry Piあたりでやりたいところですが、ひとまずMacBook Pro上でやってしまいます。(Raspberry Piでの推論はまた今度)

import numpy as np
import tensorflow as tf

interpreter = tf.lite.Interpreter(model_path="mnist.tflite")
interpreter.allocate_tensors()

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

allocate_tensors()はメモリの確保で、モデルの読み込み直後に必須らしい。

get_input_details()get_output_details()は、入力と出力の形状を示すもので、このあと使います。ちなみにこんな出力。

いざ推論です。

input_shape = input_details[0]['shape']

input_data = np.reshape(X_test[0], input_shape)

interpreter.set_tensor(input_details[0]['index'], np.array(input_data, dtype=np.float32))

interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

出力結果は、このとおり。

y_test[0]で正答が7であることが分かりますが、推論の結果を見てみると、ちょうど7にあたるところが9.9999678e-01になっています。よーよーご名答!

ということで、実際のエッジデバイスでの検証は次の機会に。

この記事を書いた人

井上 研一

株式会社ビビンコ代表取締役、ITエンジニア/経済産業省推進資格ITコーディネータ。AI・IoTに強いITコーディネータとして活動。2018年、株式会社ビビンコを北九州市に創業。IoTソリューションの開発・導入や、画像認識モデルを活用したアプリの開発などを行っている。近著に「使ってわかった AWSのAI」、「ワトソンで体感する人工知能」。日本全国でセミナー・研修講師としての登壇も多数。