TensorFlowで学習したEfficientNetモデルをTensorFlowLiteに変換する

はじめに

タイトルそのまま、TensorFlowで学習したEfficientNetのカスタムモデルを
TensorFlowLiteで扱えるように変換します。
なお、TensorFlowでの学習は本記事では扱わず、学習済みモデルの変換手順のみを記載します。

TensorFlowLiteとはざっくりいうと、
TensorFlowをAndroidやiOSなどのモバイル端末で動作できるようにするライブラリです。
詳細は公式のガイドを参照ください。
今回はAndroidデバイス上でモデルを動かしたかったので、Liteに変換することにしました。

やったことはほぼ公式ガイドに記載されていることですが、
いくつか躓きポイントがあったので、そちらも合わせてご紹介します。

また、今回はEfficientNetのモデルを対象としていますが、
TensorFlowで学習したモデルであれば、おそらく同じ手順で変換出来ます（未検証）。

変換手順

今回のモデルは

~ / assests
  / variables
  / keras_metadata.pb
  / saved_model.pb

で出力されるSavedModel形式だったため、以下の公式ドキュメントの
「SavedModelを変換する」の項を参考にしました。

実行環境

実行はGoogleColaboratory上で行いました。
実行時の主なライブラリのバージョンは以下です。

• tensorflow　2.8.2
• tflite-support 0.4.1

変換

TensorFlowの学習済みモデルは用意されている前提で進めます。

コラボ上であれば特に追加のライブラリは必要なく、以下のコードを実行して変換。

import tensorflow as tf

# モデルを変換
converter = tf.lite.TFLiteConverter.from_saved_model("saved/model/path")
tflite_model = converter.convert()

# 保存
open('savedmodel-model.tflite', 'wb').write(tflite_model)

正常に実行できると、openメソッドで指定したファイル名でtflite用モデルが出力されます。

最初、これで変換できたと思い意気揚々とAndroid上で推論を実行しようとしたのですが、
メタデータが足りないとエラーが出て実行できず。
この時点で推論ができた場合は以降の手順は不要です。

メタデータの追加

ということで変換したモデルにメタデータを追加します。
手順は公式ガイドに記載されています。

必要ライブラリのインストール

!pip install tflite-support

作業用ディレクトリの作成

!mkdir tflite_add_metadata
%cd tflite_add_metadata
# モデルの出力先ディレクトリ
!mkdir model_with_metadata

メタデータの追加用スクリプトはtensorflowの公式リポジトリで公開されているので、
以下リンクからファイルを取得してコラボのランタイム上に保存してください。

保存したファイルの以下の箇所を学習時の情報に書き換えます。

# 省略

_MODEL_INFO = {
    "mobilenet_v1_0.75_160_quantized.tflite": # 変換したモデルファイル名を指定
        ModelSpecificInfo(
            name="MobileNetV1 image classifier", # 任意の名前
            version="v1", # 任意のバージョン
            image_width=160, # 以降は学習時の情報に合わせて変更する
            image_height=160,
            image_min=0,
            image_max=255,
            mean=[127.5],
            std=[127.5],
            num_classes=1001)
}

# 省略

私の場合はmobilenet_v1_0.75_160_quantized.tfliteの部分と、
name image_width image_height num_classesを変更しました。

実行にはクラス一覧のラベルファイルも必要になりますので、用意して実行します。

!python metadata_writer_for_image_classifier.py \
    --model_file=./savedmodel-model.tflite \ # モデルファイルのパス
    --label_file=./labels.txt \ # ラベルファイルのパス
    --export_directory=./model_with_metadata # 出力先ディレクトリ

正常に実行されるとメタデータを追加したモデルファイルsavedmodel-model.tfliteと、
メタデータの内容を確認できるjsonが出力されます。
推論に必要なのはモデルファイルのみです。

おわりに

すべて公式の手順通りに実行しただけですが、ドキュメントがわかりづらく情報を探すのに結構苦労したので、同じような方の一助になれば幸いです。

おまけ

AndroidStudioのエミュレーター上で推論を実行する例をおまけでご紹介します。
とりあえずちゃんと推論できるか確認してみたかったので、
TensorFlow公式デモアプリのモデル読み込み部分を自分のモデルに差し替えただけです。

実行環境

• macOS Big Sur 11.4
• AndroidStudio Chipmunk (2021.2.1)

手順

AndroidStudioの操作は省略します。

1. TensorFlow公式リポジトリからコードをcloneします

2. cloneしてきたコードのlite/examples/image_classification/androidをAndroidStudioのプロジェクトとして開きます

3. assetsフォルダに変換したモデルを格納します

4. ImageClassifierHelper.ktの以下の箇所を用意したモデルのファイル名に書き換えます

   val modelName =
               when (currentModel) {
                   MODEL_MOBILENETV1 -> "mobilenetv1.tflite"
                   MODEL_EFFICIENTNETV0 -> "efficientnet-lite0.tflite"
                   MODEL_EFFICIENTNETV1 -> "efficientnet-lite1.tflite"
                   MODEL_EFFICIENTNETV2 -> "efficientnet-lite2.tflite"
                   else -> "mobilenetv1.tflite"
               }
   

   ↓

   val modelName = "model.tflite" //用意したモデルのファイル名
   

5. アプリを実行する

結果として自前モデルのクラスが帰ってくれば正常に変換できています。