GCP ML EngineでBatch Predictionを行う際に気をつけたいこと

## 本記事でやること
以下の公式ドキュメントである通り、BatchでPredictionを行う際ML Engine側では分散処理が行われ入力と出力の順番が一致しません。なので、**インスタンスキー**なるものを用いて入力と出力を一意にする必要があります。

つまり、入力としてインスタンスキーを明示的に書く必要があるが、Predictionの処理の中では何もせずに予測値と共に出力させるように実装しなければならないということなのですが、
公式ドキュメントに実装する際にはどうすればいいのかは丁寧に書いてなく、また日本語の文献もたくさんあるわけではなかったので備忘録的なことを含めてqiitaに投稿します。

>AI Platform は、分散処理を使用してバッチ予測ジョブを実行します。つまり、データは仮想マシンの任意のクラスタに分散され、予測できない順序で処理されます。返される予測と入力インスタンスを照合できるようにするには、インスタンス キーを定義しておく必要があります。インスタンス キーは、すべてのインスタンスが持つ値であり、一連のデータのインスタンス間で一意です。最も簡単なキーはインデックス番号です。

https://cloud.google.com/ml-engine/docs/prediction-overview?hl=ja#instance_keys



## 対象読者
- tensrorflowまたはkerasでモデルを書いたことのある方
- tensorflowのestimatorを触ったことのある方
- GCPのML Engineを触ったことのある方
- Batch Predictionと聞いてピンと来る方

## 使用言語とフレームワーク

- Python 3.6.5
- tensorflow 1.15

## 前提
試した実装方法としては2つあります。
1. estimatorをforward_features関数でラップし、インスタンスキーとなるカラム名を文字列として引数に渡す方法（おそらくこの方法が正攻法）
2. モデルから出力される予測値とインスタンスキーをconcatして最終的な出力とする方法

今回は、MovieLensのデータセット(rating.csv)を使ってuserIdとmovieIdからratingを予測するという想定で進めます。またインデックスキーには公式ドキュメントに記載がある通りデータのインデックス番号を用います。

| userId | movieId | rating | key | 
|--------|---------|--------|-----|
|1   |   1 | 1.5  | 1
|1   |   2 | 2.5  | 2
|2   |   1 | 3.5  | 3
|2   |   2 | 4.5  | 4

## 方法1: estimatorをforward_features関数でラップする
まずは[forward_features関数]
(https://github.com/tensorflow/tensorflow/blob/r1.8/tensorflow/contrib/estimator/python/estimator/extenders.py#L143)のdescriptionを見てみましょう。
しっかりBatch Predictionのことが書いています。
（冒頭で書いた通り、出力の順序が保証されていないのでinputのkeyを使ってrejoinする必要があるとの記載もあります。）

```python
def forward_features(estimator, keys=None):
  """Forward features to predictions dictionary.

  In some cases, user wants to see some of the features in estimators prediction
  output. As an example, consider a batch prediction service: The service simply
  runs inference on the users graph and returns the results. Keys are essential
  because there is no order guarantee on the outputs so they need to be rejoined
  to the inputs via keys or transclusion of the inputs in the outputs.

```
今回書いたモデルは以下のようなシンプルなモデルです。
図からわかる通り、モデルのInputLayerのインスタンスキーはどこのノードにも繋がっていない状況です。
ただし、forward_featuresの実装の中でestimatorのoutputの中にインスタンスキーを含めるような処理がなされているのでそこで繋がるのだと思います。（恐らく[ここ](https://github.com/tensorflow/tensorflow/blob/r1.8/tensorflow/contrib/estimator/python/estimator/extenders.py#L248)）
モデルの実装の中ではkeyを明示的に書くだけで特に難しいことはしていません。

![スクリーンショット 2019-12-18 9.25.30.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/270991/592e6783-e037-4938-f5db-0a93b41d5ca5.png)

```python
MAX_MOVIES=100000
MAX_USERS=100000
DIM_EMBEDDING=10

def get_model():
    key = tf.keras.layers.Input(shape=(), name='key', dtype='int32')
    
    w_inputs = tf.keras.layers.Input(shape=(1,), dtype='int32', name='movieId')
    w = tf.keras.layers.Embedding(MAX_MOVIES, DIM_EMBEDDING, name='movie')(w_inputs)
    
    u_inputs = tf.keras.layers.Input(shape=(1,), dtype='int32', name='userId')
    u = tf.keras.layers.Embedding(MAX_USERS, DIM_EMBEDDING, name='user')(u_inputs)
    
    o = tf.keras.layers.Multiply()([w, u])
    o = tf.keras.layers.Dropout(rate=0.5)(o)
    o = tf.keras.layers.Flatten()(o)
    o = tf.keras.layers.Dense(10, activation='relu')(o)
    o = tf.keras.layers.Dense(1, name='predictions')(o)
        
    model = tf.keras.Model(inputs=[w_inputs, u_inputs, key], outputs=o)
    model.summary()
    return model
```

`forward_features`関数は、modelのコンパイル/modelからestimatorを作成するなどを行うメイン関数(`train_and_evaluate`)の中で使用しています。

```python
def train_and_evaluate(outdir, train_step):
    tf.logging.set_verbosity(v=tf.logging.INFO)
    model = get_model()
    model.compile(optimizer='adam', loss="mae", metrics=["mae"])

    estimator = tf.keras.estimator.model_to_estimator(keras_model = model,
                                                     model_dir = outdir,
                                                     config = tf.estimator.RunConfig(save_checkpoints_secs=100))

    train_spec = tf.estimator.TrainSpec(input_fn=lambda: read_dataset(filename='train.csv', 
                                                                      mode=tf.estimator.ModeKeys.TRAIN),
                                       max_steps=train_step)
    
    exporter = tf.estimator.LatestExporter(name='exporter',
                                          serving_input_receiver_fn=serving_input_fn)
    
    # ここでestimatorをforward_featuresでラップしている
    estimator = tf.contrib.estimator.forward_features(estimator, 'key')
    
    eval_spec = tf.estimator.EvalSpec(input_fn=lambda: read_dataset(filename='test.csv',
                                                                   mode=tf.estimator.ModeKeys.EVAL),
                                     exporters=exporter,
                                     steps=None,
                                     throttle_secs=15)
    
    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
```


trainingではほとんどデータを入力していないので予測値が変なことになっていますが。インスタンスキーはそのままで予測値と一緒に返って来ていることがわかります。
![スクリーンショット 2019-12-18 9.27.22.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/270991/cc6b46d7-0c7e-c133-de94-3dba64e6e0ac.png)




こちらの[issue](https://github.com/tensorflow/tensorflow/issues/34418)に記載がありますが、tensorflowの2系においてはforward_featuresが属するcontrib.estimatorは非推奨となりますので、issue内に記載があるように自分でfeature_forwardなる関数を定義する必要があるのかもしれません。



## 方法2: グラフから出力される値とインスタンスキーをconcatする

以下の図のように最後のDenseから出力される値が予測値となりますが、インスタンスキーをそのままconcatし最終的な出力とします。また、インスタンスキーはスカラー値で入って来るのでReshapeしてカタチを整形しています。
こうすることでインスタンスキーには何も処理を加えず出力することが可能になります。

![スクリーンショット 2019-12-17 13.48.02.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/270991/a7423979-1246-c88c-3a37-00fec6b2e2c6.png)

```python
MAX_MOVIES=100000
MAX_USERS=100000
DIM_EMBEDDING=10

def get_model():
    key_raw = tf.keras.layers.Input(shape=(), name='key', dtype='int32')
    key = tf.keras.layers.Reshape((1,), input_shape=(), dtype='int32')(key_raw)
    
    w_inputs = tf.keras.layers.Input(shape=(1,), dtype='int32', name='movieId')
    w = tf.keras.layers.Embedding(MAX_MOVIES, DIM_EMBEDDING, name='movie')(w_inputs)
    
    u_inputs = tf.keras.layers.Input(shape=(1,), dtype='int32', name='userId')
    u = tf.keras.layers.Embedding(MAX_USERS, DIM_EMBEDDING, name='user')(u_inputs)
    
    o = tf.keras.layers.Multiply()([w, u])
    o = tf.keras.layers.Dropout(rate=0.5)(o)
    o = tf.keras.layers.Flatten()(o)
    o = tf.keras.layers.Dense(10, activation='relu')(o)
    o = tf.keras.layers.Dense(1, name='predictions')(o)
    
    # ここでpredictionとインスタンスキーをconcatしている
    # インスタンスキーはint型なので型を合わせるために無理やりfloatにしている
    pred = tf.keras.layers.Concatenate()([o, tf.cast(key, tf.float32)])
    
    model = tf.keras.Model(inputs=[w_inputs, u_inputs, key_raw], outputs=pred)
    model.summary()
    return model
```
ただし、この方法の場合Lossを計算する際にインスタンスキーを含めてしまうのでカスタム損失関数を書く必要があります。以下のようにpredictionの値だけを関数内で取得するようにすればいいだけなのでそこまで難しくはないと思います。

```python
def my_rmse(y_true, y_pred):
    return tf.keras.backend.sqrt(tf.keras.backend.mean(
        tf.keras.backend.square(y_pred[1][0] - y_true))) # y_pred[1][0]でpredictionのみを取得するようにする
```

その他の実装に関しては特に変わったことはしていません。
以下のように実際にpredictionをしてみるとこのように出力されます。

![スクリーンショット 2019-12-18 9.02.56.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/270991/8f7f0dd8-315e-9398-0dad-af8756581d63.png)

## まとめ

今回2つの方法で試してみましたが、違いとしてはtensorflow側で用意されている関数（forward_features）を使うか自力でグラフの構造を書くかくらいで処理の内容としてはほとんど変わりがないと思います。（もちろん前者の方をオススメします）
ML Engine上でのBatch Predictionに関してあまり日本語の文献がなかったので同じことで困っている方の助けになれば幸いです。