Object Detection API の SavedModel に ML Engine 用の signature を追加する

これは [TensorFlow Advent Calendar 2017](https://qiita.com/advent-calendar/2017/tensorflow) の 3 日目の記事です。

この記事では TensorFlow Object Detection API を ML Engine で利用するために、若干の改造が必要だったことについて述べます。

# Tensorflow Object Detection API とは

2017年6月に Google が公開した、画像から物体検出を行うモデルの実装とトレーニング済みモデルの TensorFlow を用いたライブラリ/ツール群です。
[Google Open Source Blog の記事](https://opensource.googleblog.com/2017/06/supercharge-your-computer-vision-models.html)
このブログ記事にあるように、画像から人や自動車などの物体が写っている位置を矩形で検出することができます。 COCO と呼ばれるデータセットでトレーニング済みのモデルも公開されているためこれはすぐに利用することもできますし、また自分でデータセットを用意することで転移学習する方法についてもできます。
Google Cloud Big Data and Machine Learning Blog で [Cloud ML Engine を使ったトレーニング方法](https://cloud.google.com/blog/big-data/2017/06/training-an-object-detector-using-cloud-machine-learning-engine) と [Cloud ML Engine を使った利用(検出)方法](https://cloud.google.com/blog/big-data/2017/09/performing-prediction-with-tensorflow-object-detection-models-on-google-cloud-machine-learning-engine) が説明されています。

# Cloud ML Engine とは

Google Cloud Platform のサービスで、主に TensorFlow を使ったトレーニングと、SavedModel 形式で書き出したモデルによる予測 API (online prediction と batch prediction に対応)を提供しています。

今回の記事は Cloud ML Engine で batch prediction を行う場合に Object Detection API の実装では困ることがある、という話です。

# Object Detection API の問題点

Cloud ML Engine のドキュメントの [Instance keys についての項目](https://cloud.google.com/ml-engine/docs/prediction-overview#instance_keys) を読むと、batch prediction は分散処理を行うため、入力はバラバラに処理され、出力結果が入力のどの画像データと対応しているのかがわからなくなってしまうので、対応付けをするために入出力に instance keys を追加しましょう、ということが書いてあります。

つまり、適当な例で書くと

> { "inputs": "データA" }
> { "inputs": "データB" }
> ...

のような入力(改行区切りのJSON)に対して、batch prediction の結果のファイルは

> { "outputs": 1.0 }
> { "outputs": 0.0 }
> ...

のようになって、どれが `"データA"` の行に対応する結果なのかわからなくなってしまうから

> { "inputs": "データA", "key": 0 }
> { "inputs": "データB", "key": 1 }
> ...

のようにユニークなキーになるフィールドを追加して

> { "outputs": 1.0, "key": 1 }
> { "outputs": 0.0, "key": 0 }
> ...

と出力に `key` をそのまま出力させるようにして、対応がわかるようにしましょうね、ということです。 [^1]

ところが Object Detection API が作成する SavedModel ではこの instance key として使える入出力が存在しません。
このままでは、[Blog Post では ML Engine で batch prediction できると書かれている](https://cloud.google.com/blog/big-data/2017/09/performing-prediction-with-tensorflow-object-detection-models-on-google-cloud-machine-learning-engine#run-batch-prediction)のに、実際には使いものにならないという残念な状態なのです。

# TensorFlow Object Detection API を修正する

幸いなことに、TensorFlow Object Detection API のコードは GitHub で公開されています。

https://github.com/tensorflow/models/tree/master/research/object_detection

そして、トレーニングやモデル構築のコードとは別に SavedModel に書きだすための部分は独立したツール/モジュールとして作られているので、この部分にだけちょっと手を入れることで予測に使う SavedModel に instance key のためのフィールドを追加することができます。

こんな感じのパッチを当てれば `key` という入力のフィールドがそのまま `key` という出力のフィールドに出力されるようになります。

```py
--- a/research/object_detection/exporter.py
+++ b/research/object_detection/exporter.py
@@ -284,11 +284,15 @@ def _write_saved_model(saved_model_path,

       builder = tf.saved_model.builder.SavedModelBuilder(saved_model_path)

+      key_placeholder = tf.placeholder(dtype=tf.int32, shape=(None,))
+
       tensor_info_inputs = {
+          'key': tf.saved_model.utils.build_tensor_info(key_placeholder),
           'inputs': tf.saved_model.utils.build_tensor_info(inputs)}
       tensor_info_outputs = {}
       for k, v in outputs.items():
         tensor_info_outputs[k] = tf.saved_model.utils.build_tensor_info(v)
+      tensor_info_outputs['key'] = tf.saved_model.utils.build_tensor_info(tf.identity(key_placeholder))

       detection_signature = (
           tf.saved_model.signature_def_utils.build_signature_def(
```

このパッチを当てたあとで [このセクション](https://cloud.google.com/blog/big-data/2017/09/performing-prediction-with-tensorflow-object-detection-models-on-google-cloud-machine-learning-engine#export-to-the-savedmodel-model) の object_detection/export_inference_graph.py を実行すれば、入出力に `key` を持つ SavedModel が export されるので、あとは同じ手順で入力ファイルにユニークな `key` フィールドに連番を追加しておくだけです。

ちなみに、同様の変更を既存の挙動を壊さないようにオプションとして追加している [pull request](https://github.com/tensorflow/models/pull/2941) が既に作られていました。 これがマージされればパッチを当てなくても良くなりそうです[^2]。

# その他の問題

今回は prediction のための記事だったので直接は関係ないのですが、実は TensorFlow Object Detection API の現在(12/2時点)の master には Cloud ML Engine でトレーニングする時に matplotlib への依存関係が存在していてエラーが発生するという状態です。[issue](https://github.com/tensorflow/models/issues/2739) が登録されていて、いくつか回避方法[^3]も提示されていますが未検証です。

TensorFlow Object Detection API の [README.md](https://github.com/tensorflow/models/blob/master/research/object_detection/README.md#release-information) には一応 Release Information というセクションがあって、リリース時におおまかな変更点は記述されているのですが、リリースといっても tag などは打たれていないのであるリリース時点のソースコードを取り出すのが簡単ではありません。軽く調べてみたところ [October 31, 2017](https://github.com/tensorflow/models/blob/master/research/object_detection/README.md#october-31-2017)のあたりの時点までは Cloud ML Engine でトレーニング可能だったようなので、そのあたりの checksum でチェックアウトしてパッケージを作成すれば使えると思います。tag を打ってくれるようになると嬉しいですね。

# まとめ

- Cloud ML Engine で batch prediction をするには、入出力の対応を取る instance key が signature に必要です
- TensorFlow Object Detection API の epxport する SavedModel は instance key を持っていないため Cloud ML Engine で batch prediction する時に問題があります
- TensorFlow Object Detection API に instance key を追加する pull request が存在します
- 現在の master では Cloud ML Engine 上でトレーニングに失敗する問題があるようです。

## 脚注

[^1]: この例では instance key の必要性を強調するためわざと結果の順番をいれかえています。また ML Engine の batch prediction ではそもそも結果が複数のファイルに分割されることがあります。
[^2]: ただしこの pull request を適用した場合は、instance key を追加するためには object_detection/export_inference_graph.py の実行時に `--instance_key_type` というオプションを指定してあげる必要があるようです。
[^3]: Tensorflow 1.4 を利用するようにパッケージ指定する、setup.py を修正して依存関係を追加する。