Including BEAM preprocessing graph in Keras models at serving

Question

Short Question:

Since Tensorflow is moving towards Keras and away from Estimators, how can we incorporate our preprocessing pipelines e.g. using tf.Transform and build_serving_input_fn() (which are used for estimators), into our tf.keras models?

---

From my understanding, the only way to incorporate this preprocessing graph is to first build the model using Keras. Train it. Then export it as an estimator using tf.keras.estimator.model_to_estimator. Then create a serving_input_fn and export the estimator as a saved model, along with this serving_input_fn to be used at serving time.

To me it seems tedious, and not the correct way of doing things. Instead, I would like to go directly from Keras to Saved Model.

---

Problem

I would like to be able to include APAHCE BEAM preprocessing graph in a Keras saved model.

I would like to serve the trained Keras model, hence I export it using SavedModel. Given a trained model, I would like to apply the following logic to predict the outcome.

raw_features = { 'feature_col_name': [100] } # features to transform
working_dir = 'gs://path_to_transform_fn' 

# transform features
transformed_features = tf_transform_output.transform_raw_features(raw_features)

model = load_model('model.h5')
model.predict(x=transformed_features)

When I define my model, I use Functional API and the model has the following inputs:

for i in numerical_features:
    num_inputs.append(tf.keras.layers.Input(shape=(1,), name=i))

This is the problem, because tensors are not inputed directly into keras from tf.Dataset, but instead, are linked using the Input() layer.

When I export the model using tf.contrib.saved_model.save_keras_model(model=model, saved_model_path=saved_model_path), I can readily serve predictions if I handle the preprocessing in a separate script. The output of this would look like

Is this what generally happens? For example, I would preprocess the features as part of some external script, and then send transformed_features to the model for prediction.

Ideally, it would all happen within the Keras model/part of a single graph. Currently it seems that I'm using the output of one graph as an input into another graph. Instead I would like to be able to use a single graph.

If using Estimators, we can build a serving_input_fn() which can be included as an argument to the estimator, which allows us to incorporate preprocessing logic into the graph.

I would like to also hear your Keras + SavedModel + Preprocessing ideas on serving models using Cloud ML

Answer 1

For your question on incorporating Apache Beam into the input function for a tf.transform pipeline, see this TF tutorial that explains how to do that:

"https://www.tensorflow.org/tfx/transform/get_started#apache_beam_implementation

On using TF 2.0 SavedModel with Keras, this notebook tutorial demonstrates how to do that:

https://www.tensorflow.org/beta/guide/keras/saving_and_serializing#export_to_savedmodel