めもめも

このブログに記載の内容は個人の見解であり、必ずしも所属組織の立場、戦略、意見を代表するものではありません。

TensorFlowのSavedModelに関するメモ

何の話かというと

TF1.0から、新しいモデルの保存形式として、SavedModelが導入されました。

このフォーマットの特徴と使い方のTipsをまとめます。

SavedModel形式で保存する方法

典型的には、Experiment APIを利用します。Experiment APIは、tf.contrib.learnの中でも最もハイレベルなAPIで、

・モデル
・トレーニング用データ入力関数
・評価用データ入力関数

を与えて実行すると、トレーニング用データでトレーニングを行いつつ、適当なタイミングで評価用データによる評価結果も出力してくれるというものです。しかも、Experiment APIを利用したコードを Cloud MLE のジョブに投げると、何も考えなくても、自動的に分散学習処理が行われます。また、学習途中のチェックポイントを自動保存するので、たとえば、global step=1000 まで実行した後、再度、同じジョブを投げると、global step=1000 のところから学習を再開してくれます。

Experiment APIを用いたコードのサンプルは、この辺りを参考にしてください。

SavedModelに関連するポイントとしては、次のボイラープレートで、serving_input_fn() を定義して、

def serving_input_fn():
    feature_placeholders = {
      'is_male': tf.placeholder(tf.string, [None]),
      'mother_age': tf.placeholder(tf.float32, [None]),
      'mother_race': tf.placeholder(tf.string, [None]),
      'plurality': tf.placeholder(tf.float32, [None]),
      'gestation_weeks': tf.placeholder(tf.float32, [None]),
      'mother_married': tf.placeholder(tf.string, [None]),
      'cigarette_use': tf.placeholder(tf.string, [None]),
      'alcohol_use': tf.placeholder(tf.string, [None])
    }
    features = {
      key: tf.expand_dims(tensor, -1)
      for key, tensor in feature_placeholders.items()
    }
    return tflearn.utils.input_fn_utils.InputFnOps(
      features,
      None,
      feature_placeholders)

Experimentオブジェクトを生成する時に、下記の形で、export_strategiesオプションにこの関数を渡します。

tflearn.Experiment(
...
        export_strategies=[saved_model_export_utils.make_export_strategy(
            serving_input_fn,
            default_output_alternative_key=None,
            exports_to_keep=1
        )]
    )

これは、トレーニングが終わったモデルをつかってPredictionを行う際の入力変数を指定するものです。トレーニングが終了すると、これらをPlaceholderで受けて、Prediction結果を出力するためのグラフが自動的に構築されて、SavedModelの一部として保存されます。これにより、Cloud MLE上で、Prediction用のAPIを公開した際の入力フォーマットが決まります。上記の例であれば、Prediction処理は次のようなコードになります。

from googleapiclient import discovery
from oauth2client.client import GoogleCredentials
credentials = GoogleCredentials.get_application_default()
api = discovery.build('ml', 'v1', credentials=credentials)

request_data = {'instances':
  [
      {
        'is_male': 'True',
        'mother_age': 26.0,
        'mother_race': 'Asian Indian',
        'plurality': 1.0,
        'gestation_weeks': 39,
        'mother_married': 'True',
        'cigarette_use': 'False',
        'alcohol_use': 'False'
      },
      {
        'is_male': 'False',
        'mother_age': 29.0,
        'mother_race': 'Asian Indian',
        'plurality': 1.0,
        'gestation_weeks': 38,
        'mother_married': 'True',
        'cigarette_use': 'False',
        'alcohol_use': 'False'
      }
  ]
}

parent = 'projects/%s/models/%s/versions/%s' % (PROJECT, 'babyweight', 'v1')
response = api.projects().predict(body=request_data, name=parent).execute()
print "response={0}".format(response)

------
response={u'predictions': [{u'outputs': 7.265425682067871}, {u'outputs': 6.78857421875}]}

細かすぎて伝わらない注意点(その1)

先ほどの served_fn() の定義内で、Placeholder を [None] のサイズで定義した後に、下記の tf.expand_dims() で [None,1] のサイズに拡張しています。謎ですね。

    features = {
      key: tf.expand_dims(tensor, -1)
      for key, tensor in feature_placeholders.items()
    }

一般的には、Placeholder のサイズは [None, x] で、この場合、この Feature は x 要素のベクトルとして与える必要があります。ただし、このサンプルでは、各 Feature はスカラーなので、仮に、Placeholder 自体を [None, 1] で定義すると、スカラー値を ['hoge'] のように1要素のベクトルとして渡す必要があって面倒です。そこで、Placeholder 自体は [None] にして、スカラー値をそのまま受け取れるようにしておき、その後で [None, 1] に拡張してから後段のグラフに渡しています。(後段のグラフは、Placeholderからの入力は、[None, 1] を期待しているので。)ああわかりにくい。

ちなみに、この次に出てくる例(MNIST)では、Feature は、784個の要素のベクトルなので、普通に [None, 784] で定義して、そのまま使っています。

細かすぎて伝わらない注意点(その2)

serving_input_fn() を定義する際に、feature_placeholdersに指定したキーがPredictionのAPIを呼ぶ際の入力項目のキーになることが分かります。ところが、むっちゃ細かい話なんですが、次のように入力項目が1つだけの場合、キーの名前が強制的に 'inputs' に変換されます。

def serving_input_fn():
  feature_placeholders = {'image': tf.placeholder(tf.float32, [None, 784]}
  features = {
    key: tensor
    for key, tensor in feature_placeholders.items()
  }
  return learn.utils.input_fn_utils.InputFnOps(
    features,
    None,
    feature_placeholders
  )

犯人はこのあたりなんですが、まったくもって余計なお世話ですね・・・。(このせいで3時間ぐらい悩みました。)

SavedModel形式の中身

SavedModel形式で保存すると、出力ディレクトリ直下に、モデルのメタ情報をprotobufで記載したファイルsaved_model.pbが保存されます。Variablesの値は、これとは別に、variablesディレクトリの下に独自のバイナリ形式で保存されます。Cloud MLE上でPrediction用APIを立ち上げる際は、出力ディレクトリを指定すると、その中のファイルを適当にあさって、必要な情報を吸い上げてくれます。

ただし、自分のコードからモデルをリストアする際は、メタ情報の内容を理解しておく必要があります。export_strategiesを指定する際、次のように「as_text=True」を指定すると、テキスト形式のsaved_model.pbtxtが出力されるので、このファイルを開くと、メタ情報の構造が確認できます。

tflearn.Experiment(
...
        export_strategies=[saved_model_export_utils.make_export_strategy(
            serving_input_fn,
            default_output_alternative_key=None,
            exports_to_keep=1,
            as_text=True
        )]
    )

また、自前のPythonコードからSavedModel形式のモデルをロードする関数が、tf.saved_model.loader.load() として用意されています。これを利用して自前のコード上でPredictionを行う例がこちらになります。メタ情報から、入出力項目のキーに対応するTensorオブジェクトの名前を取得するという一手間を加えてあります。

import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

export_dir = 'モデルを保存したディレクトリ'

def run_inference_on_image():
    mnist = input_data.read_data_sets("/tmp/data/", one_hot=True)
    
    with tf.Session(graph=tf.Graph()) as sess:
        meta_graph = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], export_dir)
        
        model_signature = meta_graph.signature_def['serving_default']
        input_signature = model_signature.inputs
        output_signature = model_signature.outputs
        
        input_tensor_name = input_signature['inputs'].name
        prob_tensor_name = output_signature['probabilities'].name
        label_tensor_name = output_signature['classes'].name

        images = sess.graph.get_tensor_by_name(input_tensor_name)
        prob = sess.graph.get_tensor_by_name(prob_tensor_name)
        label = sess.graph.get_tensor_by_name(label_tensor_name)
        label_pred, prob_pred = sess.run([label, prob], feed_dict={images: mnist.test.images[0:10]})
        return label_pred, prob_pred

ちなみに、これは、下記のサンプルコードで学習したモデルをリストアして利用する例です。

Cloud MLE and GCE compatible TensorFlow distributed training example

このモデルでは、serving_input_fn()で定義した入力項目が1つなので、入力項目のキーが強制的に 'inputs' になっています。また、このモデルは、tf.contrib.learnのカスタムEstimatorを使用しており、出力項目のキーについては、下記のようにカスタムEstimator作成のお作法に従って指定が行われています。

...
  predictions = {
      "classes": tf.argmax(input=logits, axis=1),
      "probabilities": tf.nn.softmax(logits)
  }

  return model_fn_lib.ModelFnOps(mode=mode, loss=loss, train_op=train_op,
                                 predictions=predictions)