Things to keep in mind when doing Batch Prediction on GCP ML Engine

What to do in this article

As shown in the official document below, when Prediction is performed in Batch, distributed processing is performed on the ML Engine side and the order of input and output does not match. So you need to use something called a ** instance key ** to make the inputs and outputs unique.

In other words, it is necessary to explicitly write the instance key as input, but it must be implemented so that it is output with the predicted value without doing anything in the Prediction process. I didn't write carefully what to do when implementing it in the official document, and there weren't many Japanese documents, so I'll post it on qiita, including a memorandum.

AI Platform uses distributed processing to run batch prediction jobs. That is, the data is distributed across any cluster of virtual machines and processed in an unpredictable order. You must have an instance key defined so that you can match the returned predictions with the input instance. An instance key is a value that every instance has and is unique across a set of data instances. The simplest key is the index number.

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

Target audience

--Those who have written models with tensrorflow or keras --Those who have touched the tensorflow estimator --Those who have touched GCP's ML Engine --Those who hear "Batch Prediction" and come to mind

Language and framework

• Python 3.6.5
• tensorflow 1.15

Premise

There are two implementation methods I have tried.

1. Wrap the estimator with the forward_features function and pass the column name that will be the instance key as a character string as an argument (probably this method is a straightforward method)
2. How to concat the predicted value and instance key output from the model to make the final output

This time, we will proceed on the assumption that the rating is predicted from the userId and movieId using the MovieLens dataset (rating.csv). For the index key, use the index number of the data as described in the official document.

Method 1: Wrap the estimator with the forward_features function

First, [forward_features function] Take a look at the description (https://github.com/tensorflow/tensorflow/blob/r1.8/tensorflow/contrib/estimator/python/estimator/extenders.py#L143). I'm writing about Batch Prediction. (As I wrote at the beginning, there is also a statement that the output order is not guaranteed, so it is necessary to rejoin using the input key.)

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.

The model I wrote this time is a simple model as follows. As you can see from the figure, the instance key of the model InputLayer is not connected to any node. However, in the implementation of forward_features, the process of including the instance key in the output of the estimator is done, so I think that it will be connected there. (Probably here) In the model implementation, just writing the key explicitly is not particularly difficult.

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

The forward_features function is used in the main function (train_and_evaluate) that compiles the model / creates the estimator from the model.

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)
    
    #Forward the estimator here_Wrapped in 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)

In training, I haven't entered much data, so the predicted values are strange. You can see that the instance key is returned as it is with the predicted value.

There is a description in this issue, but in tensorflow 2 system, contrib.estimator to which forward_features belongs is deprecated, so in the issue You may need to define your own function called feature_forward as described.

Method 2: Concat the value output from the graph and the instance key

As shown in the figure below, the value output from the last Dense is the predicted value, but the instance key is concated as it is and used as the final output. Also, since the instance key comes in as a scalar value, it is reshaped to shape it. By doing this, it is possible to output the instance key without any processing.

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)
    
    #Here we are concating the prediction and the instance key
    #Since the instance key is an int type, it is forced to float to match the type.
    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

However, with this method, you need to write a custom loss function because the instance key is included when calculating Loss. I don't think it's that difficult because you only have to get the value of prediction in the function as shown below.

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]To get only prediction with

Nothing has changed about the other implementations. If you actually try prediction as shown below, it will be output like this.

Summary

I tried two methods this time, but the difference is that there is almost no difference in the processing content depending on whether you use the function (forward_features) prepared on the tensorflow side or write the graph structure by yourself. I will. (Of course, the former is recommended) I didn't have much Japanese literature on Batch Prediction on ML Engine, so I hope it helps those who are having trouble with the same thing.