SavedModelについてのまとめ

この記事は[TensorFlow Advent Calendar 2018](https://qiita.com/advent-calendar/2018/tensorflow)の19日目の記事です。最近お仕事で使う機会が多く、いろいろと調べたのでSavedModelについてまとめようと思います。

# SavedModelとは？
SavedModelはTensorFlowでモデルを保存する際の汎用的なシリアライゼーション形式です。
TensorFlowには基本的なモデルの保存型式としてCheckpointがあります。Pythonでコードを書いてモデルを訓練して、一定間隔毎にモデルを保存する際はCheckpoint形式の保存が適しています。しかし、出来上がったモデルをいざ本番環境で推論させたいというような場合を考えると、その推論用のプログラムを別で書く必要があったり、そもそもモデルを搭載するシステムがPythonで書かれていなかったりと、Checkpoint形式だと都合が悪い時があります。
SavedModel形式の保存を行うと言語依存せずに、コマンドラインで推論を行ったり、TensorFlow ServingやTensorFlow Model Analysis、GCPのCloud ML Engineなどの高レベルなサービスやツールでTensorFlowのモデルを利用することができるようになります。

# 保存方法

まずはモデルとして超シンプルな計算グラフを用意します。

```python:model.py
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による保存

```python:model.py
tf.saved_model.simple_save(sess, './models', inputs={"input": x}, outputs={"output": out})
```
これを実行すると`models`ディレクトリ直下に以下のようなファイルが現れます。<img width="739" alt="スクリーンショット 2018-12-18 15.38.53.png" src="https://qiita-image-store.s3.amazonaws.com/0/57774/809f1ed0-f7c5-71ae-81b8-69063f69f277.png">
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()`を用いた保存方法を紹介します。

```python:model.py
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を定義しています。各種メソッドの違いなどを知りたい場合は[ここ](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/g3doc/signature_defs.md)を参考にするといいと思います。また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の構造は以下のようになります。

```text
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なんて使わないぞ :rage:  というご意見も有るとおもいますし、TF2.0以降ではKerasとよりタイトに統合していきKerasAPI（`tf.keras`）を推奨していくという流れらしいので`tf.keras`からSavedModelを生成するための方法を参考までに記載したいと思います。

```python
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で使ったりという記事も車内でまとめてあったりするので逐次公開していこうと思います。

#  参考サイト
- [TensorFlow公式ドキュメント](https://www.tensorflow.org/guide/saved_model)
- [TensorFlow: How and why to use SavedModel](https://stackoverflow.com/questions/46513923/tensorflow-how-and-why-to-use-savedmodel)
- [Training and Serving ML models with tf.keras](https://medium.com/tensorflow/training-and-serving-ml-models-with-tf-keras-fd975cc0fa27)
- [Standardizing on Keras: Guidance on High-level APIs in TensorFlow 2.0](https://medium.com/tensorflow/standardizing-on-keras-guidance-on-high-level-apis-in-tensorflow-2-0-bad2b04c819a)