Inception-v3 を fine-tuning したモデルを tensorflow-lite で使う

はじめに

普段研究ではPyTorchを使うことが多かったのですが、勉強のために最近Tensorflow2.xに手を出しています。
またiOS開発(Swift)経験もほぼはじめまして状態で手探りの中やっていきます。
さらに研究では自然言語処理をメインでやっていましたが、画像はほぼハジメマシテです。
対戦よろしくおねがいします。

目標

• iOSで画像(鳥の種類)の分類がしたい
• DeepLearning的なアプローチがいい
• サーバーは使いたくない(お金とか色々めんどくさい)
  • iOS上で動くらしいTensorflow Liteを利用
  • 複雑なモデルの利用は難しい？
  • やってみないとわからない

んじゃやってみよう

実装

まずはデータ集め

鳥の画像データは Kaggle にそれっぽいものがあったのでもらってきます。

250 Bird Species | Kaggle

URLは100-bird-speciesって書いてありますが、なんか250種類に増えています。
つまり250クラスの分類になりますね、多いな・・・

方針

今回はサクッとよさげなものが作りたいお気持ちが強いのでつよつよモデルを組んだりとか、複雑なモデルを組むのではなく、既存の学習済みモデルをfine-tuningする形でモデルを作成します。
とくに複雑すぎるモデルを作るとtf-liteに変換するときに詰むことがあるらしいです。

よって

• 良さげな事前学習モデルを探す
• データを使って学習する
• tf-liteで扱えるモデルへ変換する
• iOSで動かしてみる

の順番に進めていきます。
ここではtf-liteで扱えるモデルに変換して推論を行うところまでやります。
iOSにはまだ乗せません。

事前学習モデルさがし

いろいろ触ってみた感じで Inception-v3 がなんか良さげだったのでこちらでやってみます。
他にもResnetとかmobilenetみたいなのもあるらしいです。軽さで言えばmobilenetがベターかもしれないのですが、あまり精度が出なかったため見送りました。(Mobilenet-v3は事前学習済みの重みがKerasでは使えないらしいのが原因かもしれない？)

学習する

モデル作成・学習の方針は以下の感じ。
画像はじめましてなので見様見真似です。

• 画像は224*224にリサイズする
• Inception-v3 の出口付近はパラメータのフリーズを解除する
• Batch Normalizationは場所関係なくフリーズを解除する
• Data Augmentationをする
  • 水平反転
  • 最大15°回転
  • 最大10%の拡大
• Validationに対してAccuracyを算出してEarlyStoppingをする

JupyterLabでやってるのでimport等は省略します。
拾い物のコードを改造しているのでぐちゃぐちゃですがご容赦を
参考: https://www.kaggle.com/eabgupt/birdclassification https://ossan-ml.hateblo.jp/entry/2019/03/23/221710

# import 省略
train_path = "./dataset/train"
valid_path = "./dataset/valid"
test_path = "./dataset/test"

train_datagen_augmented = ImageDataGenerator(
    rescale=1/255,
    horizontal_flip=True,    # 水平反転
    rotation_range=15,       # 回転
    zoom_range=0.1,          # 拡大
)

valid_datagen = ImageDataGenerator(rescale=1/255)
test_datagen = ImageDataGenerator(rescale=1/255)

train_generator_augmented = train_datagen_augmented.flow_from_directory(
    train_path,
    target_size=(224, 224),
    batch_size=32,
    color_mode="rgb",
    class_mode="sparse",
    shuffle=True,
)

validation_generator = valid_datagen.flow_from_directory(
    valid_path,
    target_size=(224, 224),
    batch_size=32,
    color_mode="rgb",
    class_mode="sparse")

test_generator = test_datagen.flow_from_directory(
    test_path,
    target_size=(224, 224),
    batch_size=32,
    color_mode="rgb",
    class_mode="sparse")

# ここからモデル定義とか
backend.clear_session()

# pre-trainモデルのload
incbasemodel4 = InceptionV3(weights="imagenet", include_top=False, input_shape= (224, 224, 3))
finetune_at = 250
incbasemodel4.trainable = True

# 最上層あたり以外をフリーズ
for layer in incbasemodel4.layers[:finetune_at - 1]:
    layer.trainable =  False
    # Batch Normalization の freeze解除
    if layer.name.startswith('batch_normalization'):
        layer.trainable = True

# Sequential に追加していく
model_inceptionv3_finetune = models.Sequential()
model_inceptionv3_finetune.add(incbasemodel4)
model_inceptionv3_finetune.add(GlobalAveragePooling2D())
model_inceptionv3_finetune.add(layers.Dense(1024, activation="relu"))
model_inceptionv3_finetune.add(layers.Dense(250, activation="softmax"))

model_inceptionv3_finetune.compile(optimizer=optimizers.Adam(),
               loss="sparse_categorical_crossentropy",
               metrics=["accuracy"])

checkpoint_path = "checkpoints/cp-{epoch:04d}.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)

# checkpoint保存用コールバック
cp_callback = ModelCheckpoint(
    filepath=checkpoint_path, 
    verbose=1, 
    save_weights_only=True,
    period=1
)

# TensorBoard用コールバック
tb_callback = TensorBoard(log_dir="./logs/finetune-augmented", histogram_freq=1, write_graph=True)

# fit_generatorは非推奨らしいが、動くのでひとまずこれで
history = model_inceptionv3_finetune.fit_generator(
          train_generator_augmented,
          epochs=100,
          validation_data=validation_generator,
          verbose=1, 
          shuffle=True,
          callbacks=[
              EarlyStopping(monitor="val_accuracy", patience=10, restore_best_weights=True),
              cp_callback,
              tb_callback,
          ]
)

ぐちゃぐちゃですが、こんな感じで学習させましょう。

モデルの図はこんな感じ。シンプルです。

赤がtrainで緑がvalidationです。

まあよさげですね、validationで最強のAccを示したやつを採用しましょう。
個人的にEarlyStoppingはあまり好きじゃない(露骨にvalidation lossが上がってない限り、まだいけるのでは？と思うことがある)ですが、今回のようにどういう構造で行くか模索中のときはありかな、と思いました。
実際に最強モデル作る準備ができたらEarlyStoppingは外してTensorBoardとにらめっこがいいと思います。

また、LearningRateSchedulerとかを使って学習率を動的に変動させて、さらなる高みを目指すのもありだと思います。
気力があったらやるかも。

変換する

参考: https://qiita.com/PINTO/items/008c54536fca690e0572

まずは最高の結果を示したモデルをロードしましょう。
フリーズが必要なのかというといらない気がするけど念の為できる限り重み以外同じ状況を作り上げます。

# epoch=18のもの(0から数えると17だけど)
checkpoint_file = "./checkpoints/finetune-augmented/cp-0018.ckpt"

incbasemodel4 = InceptionV3(weights="imagenet", include_top=False, input_shape=(224, 224, 3))

finetune_at = 250
incbasemodel4.trainable = True

# 最上層あたり以外をフリーズ
for layer in incbasemodel4.layers[:finetune_at - 1]:
    layer.trainable = False
    # Batch Normalization の freeze解除
    if layer.name.startswith('batch_normalization'):
        layer.trainable = True

model = models.Sequential()
model.add(incbasemodel4)
model.add(GlobalAveragePooling2D())
model.add(layers.Dense(1024, activation="relu"))
model.add(layers.Dense(num_classes, activation="softmax"))

model.compile(optimizer=optimizers.Adam(),
               loss="sparse_categorical_crossentropy",
               metrics=["accuracy"])

model.load_weights(checkpoint_file)

loss, acc = model.evaluate_generator(test_generator)
print("Test Acc:", acc)

Test Acc: 0.9879999756813049

Testでは98.80%のAccを示しました。
250クラスの分類ですのでなかなか良い結果なのではないでしょうか？
Kaggleに99.5%とかありますが、これは100-bird-speciesだったころの記録みたいです。

入出力を一応調べてsaved_modelとして保存します。

print("input:", model.inputs)
print("output:", model.outputs)
model.save("saved_model/bird_model")

input: [<KerasTensor: shape=(None, 224, 224, 3) dtype=float32 (created by layer 'inception_v3_input')>]
output: [<KerasTensor: shape=(None, 250) dtype=float32 (created by layer 'dense_1')>]

そしてtfliteファイルに変換します。

converter = tf.lite.TFLiteConverter.from_saved_model('./saved_model/bird_model/')
converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
tflite_quant_model = converter.convert()
with open('./bird_model.tflite', 'wb') as w:
    w.write(tflite_quant_model)
print("Weight Quantization complete! - bird_model.tflite")

tfliteファイルを読み込んでtf-liteで使えるかどうかを試します。

# tfliteの読み込みと推論をしてみる
lite_model = tf.lite.Interpreter(model_path="bird_model.tflite")
lite_model.allocate_tensors()

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

# generatorなのになんかアクセスできる
data = test_generator[0]

lite_model.set_tensor(input_details[0]['index'], [data[0][0]])
lite_model.invoke()
pred = int(np.argmax(lite_model.get_tensor(output_details[0]['index'])))
gold = int(data[1][0])

print("Predicted:", pred)
print("Gold:", int(data[1][i]))

Predicted: 80
Gold: 80

無事に推論できました！
量子化を行っている関係で精度が落ちているとは思います。適当にコードを書いて調べたら ACC: 0.9861751152073732 と表示されました。確かに落ちたけど、思ったよりは落ちてませんね。

ハマったところ

tfliteに変換した後、読み込めなかった

ValueError: Invalid tensors なんたらかんたら言われました。
これは Inception-v3 の input_shape を横着して指定しないとこうなります。
input_shape=(224, 224, 3) みたいに指定してあげたら読み込めました。

学習がやたら遅い

画像ほぼハジメマシテなのでこんなもんかと思っていたのですが、なんかおかしいなって思って確認。
なんと環境がおかしくなっていてGPUのメモリを占拠するだけして実際にはGPUを使えていませんでした。(nvidia-smi上はメモリを使っていることになっていたが、TensorflowからはGPUが見えていなかったみたい？)
テキトーにCUDAとか入れ直したら環境が生き返りました。

Next...

iOSで読み込み/推論が行えるか試したいと思います。
もしうまくいかなければ別のモデルを作る必要が出てきます。