Kerasで学習済のモデルをTensorflowのSavedModelに簡単に変換する

変換する理由

Google Cloud ML Engine上で予測させるため

やり方

convert.py

import argparse

from keras import backend as K
from keras.applications import mobilenetv2
from keras.layers import DepthwiseConv2D
from keras.models import load_model
import tensorflow as tf
from tensorflow import saved_model


def run(args):
  sess = K.get_session()
  model = None
  if args.base == 'mobilenet':
    model = load_model(args.model, custom_objects={
        'relu6': mobilenetv2.mobilenet_v2.relu6,
        'DepthwiseConv2D': DepthwiseConv2D})
  else:
    model = load_model(args.model)

  saved_model.simple_save(sess,
                          args.output,
                          inputs={'input': model.input},
                          outputs={'output': model.output})


if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('--model', type=str, default='model.h5')
  parser.add_argument('--base', type=str, default='mobilenet')
  parser.add_argument('--output', type=str, default='output')
  args = parser.parse_args()
  run(args)

注意するのは、keras2.2.0からrelu6がmobilenetv2に移動してて、DepthwiseConv2Dがlayersに移動している点。

あとはフォルダ内のファイルをGCSに入れてから、以下のコマンドを実行するだけ。
GCPとAPIの有効化は実施済とします。

deploy.sh

MODEL=$1
VERSION=$2
BUCKET=$3
gcloud ml-engine models create $MODEL
gcloud ml-engine versions create $VERSION --model $MODEL --origin BUCKET --runtime-version 1.8

これだけ。

備考

Google Cloud ML Engineへのリクエストサイズ制限は1.5MBまで。（GAEでもあったなこれ。。）
224x224だと軽く超えるので128x128がくらいが限界っぽい。

感想

tf.saved_model.simple_saveだけで変換できるってすごくない？？？