この記事はTensorFlow Advent Calendar 2018の19日目の記事です。最近お仕事で使う機会が多く、いろいろと調べたのでSavedModelについてまとめようと思います。
SavedModelとは？
SavedModelはTensorFlowでモデルを保存する際の汎用的なシリアライゼーション形式です。
TensorFlowには基本的なモデルの保存型式としてCheckpointがあります。Pythonでコードを書いてモデルを訓練して、一定間隔毎にモデルを保存する際はCheckpoint形式の保存が適しています。しかし、出来上がったモデルをいざ本番環境で推論させたいというような場合を考えると、その推論用のプログラムを別で書く必要があったり、そもそもモデルを搭載するシステムがPythonで書かれていなかったりと、Checkpoint形式だと都合が悪い時があります。
SavedModel形式の保存を行うと言語依存せずに、コマンドラインで推論を行ったり、TensorFlow ServingやTensorFlow Model Analysis、GCPのCloud ML Engineなどの高レベルなサービスやツールでTensorFlowのモデルを利用することができるようになります。
保存方法
まずはモデルとして超シンプルな計算グラフを用意します。
import tensorflow as tf
x = tf.placeholder(tf.float32, shape=[None], name='input' )
y = tf.placeholder(tf.float32, shape=[None], name='target')
a = tf.Variable(tf.zeros([1]), name='var_1')
b = tf.Variable(tf.zeros([1]), name='var_2')
out = tf.identity(a * x + b, name='out')
loss = tf.reduce_mean(tf.square(y-out))
global_step = tf.train.get_or_create_global_step()
train = tf.train.GradientDescentOptimizer(0.5).minimize(loss, global_step=global_step)
init = tf.global_variables_initializer()
sess = tf.InteractiveSession()
sess.run(init)
本来は訓練などを行なう必要が有りますが、簡単のために諸々省略しています。モデルの定義ができたので、まずはもっとも簡単な方法であるtf.saved_model.simple_saveメソッドを用いてSavedModel形式でモデルをエクスポートしてみます。
simple_saveによる保存
tf.saved_model.simple_save(sess, './models', inputs={"input": x}, outputs={"output": out})
これを実行するとmodelsディレクトリ直下に以下のようなファイルが現れます。
Sessionやモデルのディレクトリを引数とするところはtf.train.Saver()によるCheckpoint形式のモデル保存と同じですが異なるのが入力と出力のTensorを指定しているところです。当然ですが計算グラフを高築している側の人間からするとどのTensorが入力でどこが出力かということはわかりますが、TensorFlow Servingなどのライブラリからしたらそんなこと知ったこっちゃありません。また、例えばクラス分類をするモデルを考えた時、出力といってもソフトマック関数を作用させた各クラスである確率がアプリケーション側に必要なケースもあれば、単に入力がどのクラスに属するかだけが欲しい場合もあり、後者の場合はソフトマックス関数を作用させる前のTensorだけあれば十分です。そのため、SavedModelをエクスポートする際は入出力の情報を指定してあげる必要があり、そうした入出力の情報をTensorFlowでは"Signature"と呼びます。
この簡単な方法でエクスポートしたSavedModelはすぐにTensorFlow Servingやsaved_model_cliで利用することができます。また、Python以外の言語からも読み込むことも可能です。
SavedModelBuilderによる保存
先ほどのtf.saved_model.simple_saveは便利で使いやすいですがミニマムの機能しか提供していません。ここでは細々としたカスタマイズが可能なtf.saved_model.builder.SavedModelBuilder()を用いた保存方法を紹介します。
builder = tf.saved_model.builder.SavedModelBuilder('./models')
signature = tf.saved_model.predict_signature_def(inputs={'input': x}, outputs={'output': out})
builder.add_meta_graph_and_variables(sess=sess,
tags=[tf.saved_model.tag_constants.SERVING],
signature_def_map={'predict': signature})
builder.save()
simple_saveに比べて少しだけ複雑になっています。まずディレクト名を引数にとり、tf.saved_model.builder.SavedModelBuilder()を宣言します。こうすることでmodelsというディレクトリが出来上がりますがまだ中身は空です。次にSignatureを定義しています。ここは先程のsimple_saveと似ています。Signitureの定義に必要なものは入力と出力のTensorの他にメソッドと呼ばれるものが実は必要で、ここではPredictと呼ばれるメソッド用のSignatureを定義しています。各種メソッドの違いなどを知りたい場合はここを参考にするといいと思います。またPredict以外のメソッド用のSignatureを定義する関数や、より汎用的なSignatureを作成するためのtf.saved_model.build_signature_def()というメソッドも存在するので必要に応じて使い分けてください。Signatureを定義できたら、add_meta_graph_and_variablesメソッドを用いてモデルを追加します。SessionやSignatureを渡すのは当たり前なのですが、ここでもう一つ新しい概念であるtagというものが出てきます。
このtagは一言で言うと複数の計算グラフを区別するためのものです。実はSavedModelには複数の計算グラフを含めることができます。ユースケースとして、訓練用の計算グラフと推論用の計算グラフを分けるということはよく行われます。また、同じモデルであったとしてもデプロイするエッジデバイスによってGPUが乗っている場合、乗っていない場合など様々です。そうした際にはCPU版とGPU版でモデルを分けておくことは自然です。この例では１つしか計算グラフを追加していませんが、複数追加する場合、初回は必ずadd_meta_graph_and_variablesメソッドでモデルの骨組みとVariableの両方を加える必要がありますが、2回目以降はadd_meta_graphというメソッドで重みを共有する異なる計算グラフを追加することができます。tagは文字列のリストであるなら何でも良いのですが、よく使う項目としてTensorFlowでは
- SERVING
- TRAINING
- GPU
- TPU
といった固定の文字列が定義されています。GPUにおける推論用の場合は[SERVING, GPU]のようにtagを設定しておくとロードする際にわかりやすいでしょう。最後にsaveメソッドでSavedModelをエクスポートします。
構造
この例ではフルに機能を活用していないので出てきていない項目も有りますが、SavedModelの構造は以下のようになります。
assets/
assets.extra/
variables/
variables.data-?????-of-?????
variables.index
saved_model.pb|saved_model.pbtxt
saved_model.pbはメタグラフといって、計算グラフの構造を表すものです。どのオペレーションとどのVariableが足し算されてといった構造を保持しており、Variableの値そのものは保持していません。variables直下の二つのファイルがVariableの値を保持しています。このあたりはtf.train.Saver()を用いた保存とほとんど同じなのでわかりやすいでしょう。assetsは単語辞書ファイルなどの外部ファイルが保存されるフォルダになります。add_meta_graph_and_variables()メソッドの引数assets_collectionで指定できるようです。この辺りがCheckpoint形式と異なるところです。
tf.kerasからのエクスポート
上記の例ではTensorFlowのLowLevelAPIで記述しておりSessionなどの概念がでてきていました。Kerasとかで作るからSessionなんて使わないぞ 
というご意見も有るとおもいますし、TF2.0以降ではKerasとよりタイトに統合していきKerasAPI（tf.keras）を推奨していくという流れらしいのでtf.kerasからSavedModelを生成するための方法を参考までに記載したいと思います。
import tensorflow as tf
model = tf.keras.models.load_model('./model.h5')
export_path = './models'
with tf.keras.backend.get_session() as sess:
tf.saved_model.simple_save(
sess,
export_path,
inputs={'input': model.input},
outputs={t.name:t for t in model.outputs})
tf.keras.backend.get_session()メソッドでSessionが取得できるのであとはそのままという感じですね。ただ、直截SavedModelを吐き出す機能は見た所まだ内容なので、一手間かかりますね。TF2.0以降追加されるのを期待しましょう。
おわりに
あまりSavedModelに関する日本語の記事がなかったので書いてみました。このあとTensorFlowSerivingで使ったりという記事も車内でまとめてあったりするので逐次公開していこうと思います。