TensorFlowの学習済みモデルをKeras用に変換して使う

この記事はTensorFlow Advent Calendar 2018の10日目の記事です。

本記事では、MMdnnを使って、TensorFlowの学習済みモデルをKeras用に変換する方法を説明します。

準備

関連パッケージをインストールします。MMdnnは最新版ではやりたいことができなかったので、少し修正したものを使用します。

pip install tensorflow==1.12.0 keras==2.2.4 git+https://github.com/mktozk/MMdnn

今回はここのモデルを変換してみます。

まず、TensorFlowの学習済みモデルをダウンロードして展開します。

wget http://download.tensorflow.org/models/deeplabv3_mnv2_dm05_pascal_trainaug_2018_10_01.tar.gz
tar xf deeplabv3_mnv2_dm05_pascal_trainaug_2018_10_01.tar.gz \
  deeplabv3_mnv2_dm05_pascal_trainaug/frozen_inference_graph.pb --strip-components 1

また、後ほどスクリプトを使用するので、TensorFlowのコードをクローンしておきます。

git clone https://github.com/tensorflow/tensorflow
cd tensorflow
git checkout v1.12.0
cd ../

変換

TensorFlowの学習済みモデルを変換するためには、入出力のノード名と入力のサイズを指定する必要があります。MMdnnにはここに書かれているようにmmvismetaという可視化ツールがあるのですが、これを使うためには.metaファイルが必要です。今回は.metaファイルが手に入らなかったので、以下のようにしてTensorBoard上でノード名を確認しました。

ここではimport_pb_to_tensorboard.pyを使用します。これは、tf.GraphDefからTensorBoard用のファイルを出力してくれるスクリプトです。これを実行した後、TensorBoardを起動します。

python tensorflow/tensorflow/python/tools/import_pb_to_tensorboard.py \
  --model_dir frozen_inference_graph.pb --log_dir logs

tensorboard --logdir=logs

ブラウザでlocalhost:6006にアクセスすると、以下のような図が確認できます。

今回、入力ノードをMobilenetV2/MobilenetV2/input、出力ノードをlogits/semantic/BiasAddとしました。

最後に、以下のスクリプトを実行するとKerasのモデルが得られます。

convert_pretrained_model.sh

#!/bin/bash

INPUT_NAMES=MobilenetV2/MobilenetV2/input
OUTPUT_NAMES=logits/semantic/BiasAdd
TMP_MODEL=tmp_model

python tensorflow/tensorflow/python/tools/optimize_for_inference.py \
  --input=frozen_inference_graph.pb --output=optimized_graph.pb \
  --frozen_graph=True --input_names=${INPUT_NAMES} --output_names=${OUTPUT_NAMES}

mmtoir -f tensorflow -iw optimized_graph.pb \
  --inNodeName ${INPUT_NAMES} --inputShape 513,513,3 \
  --dstNodeName ${OUTPUT_NAMES} -o ${TMP_MODEL}

mmtocode -f keras -n ${TMP_MODEL}.pb -iw ${TMP_MODEL}.npy -o base_keras_model.py

mmtomodel -f keras -in base_keras_model.py -iw ${TMP_MODEL}.npy -o keras_model.h5

mmconvertコマンドでもKerasのモデルへ変換可能ですが、後述のようにモデル読み込み時に中間コードを使用したいため、このようにstep-by-stepで実行しています。

変換したモデルの読み込み方

変換したモデルは以下のようにして使うことが可能です。resize_bilinearはKerasのオブジェクトではないので、モデルを読み込む際に渡してあげる必要があります。

model.py

import keras

from base_keras_model import resize_bilinear


def build_model():
    model = keras.models.load_model(
        'keras_model.h5', custom_objects={'resize_bilinear': resize_bilinear})
    return model

おわりに

巨人の肩の上に立って楽をしていきましょう。