In [ ]:
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Recently, someone asked me a problem related to doing distributed prediction on a batch prediction request, and returning the results in real-time. The problem they had is that the backend of their application which sends the batch request does not know what order the individual predictions will arrive in, and thus which prediction result when with which prediction request.
In this case, they had an existing trained tf.keras model and wanted to retrofit the model such that each individual prediction request had a unique identifier and the prediction result includes the corresponding unique identifier; which we will refer herein as the example key.
This is so easy in Keras using the Functional API with multiple inputs and multiple outputs, and using the trained Keras model object as a callable (i.e., a layer). Here's the basic steps we will follow:
1. Build and train the existing model.
2. Make a second model (we call the wrapper model), with
A. Two inputs
B. Two outputs
3. The two inputs to the wrapper model are:
A. The input layer of the trained model
B. A 1D vector for the example key
4. The two outputs to the wrapper model are:
A. The output from the trained model
B. The value (example key) of the 1D vector.
That's all you need. Let's show an example of doing this with a tf.Keras model for MNIST.
In [ ]:
import tensorflow as tf
from tensorflow.keras import Input, Model
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
import numpy as np
# Reguired if you are using pre TF-2.0
# i.e, eager execution is required for adding example key as a graph operation
tf.enable_eager_execution()
# Display what version of TF you are using
print(tf.__version__)
In [ ]:
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = (x_train / 255.0).astype(np.float32)
x_test = (x_test / 255.0).astype(np.float32)
y_train = to_categorical(y_train)
y_test = to_categorical(y_test)
print("Train", x_train.shape, y_train.shape)
print("Test ", x_test.shape, y_test.shape)
We will use the Functional API to build essential a sequential dense neural network (DNN), as follows:
1. Flattens the MNIST images (28x28) into a 1D vector (784).
2. An input and hidden dense layer of 128 nodes with a ReLU activation function.
3. An output layer of 10 nodes (1 per digit) with a softmax activation function.
In [ ]:
inputs = Input((28, 28))
x = Flatten()(inputs)
x = Dense(128, activation='relu')(x)
x = Dense(128, activation='relu')(x)
outputs = Dense(10, activation='softmax')(x)
model = Model(inputs, outputs)
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['acc'])
Let's do a quick training of this MNIST model for 5 epochs. We should see ~98% accuracy on the training data. Let's not worry about the test data or overfitting. We just need a quick model to demonstrate how to make a wrapper model using the tf.Keras Functional API and pass through the example key during prediction.
In [ ]:
#Train the model
model.fit(x_train, y_train, epochs=5)
In [ ]:
model.summary()
If you have not used multiple inputs and multiple outputs, this maybe the easiest example you will see. In the Functional API, when we build the model using the Model() class, we pass two parameters, the input tensor and the output layer; which I call pulling it all together connecting the inputs to the outputs:
my_model = Model(inputs, outputs)
I won't go into detail about how to (and why) to create models with more than one input and/or output. We will simply use this method to implement passing through an example key at prediction with an existing trained tf.keras model (i.e., MNIST model). The syntax for specifying both multiple inputs and outputs looks like this:
my_model = Model( [inputs1, inputs2], [outputs1, outputs2])Let's get started. We can do this in two lines of Keras code!
1) First, we need to define a second input for passing through the example key. It will need to be number. For prediction purposes, we will define this as a scalar value. We will name this input tensor as key.
key = Input(1)
2) Next. we build a new model ("the wrapper") reusing the existing trained model, but we are going to add the key tensor to both the inputs and outputs. To resuse the trained model, we need to get the trained model's input tensor and output layer, as in:
Model(model.input, model.output)
Let's now expand on this and add in the key:
wrapper_model = Model( [model.input, key], (model.output, key])
That's it. We did it. We don't need to compile or train --because we reused the tf.keras model that is already trained!
In [ ]:
# Create the second input to the model for passing through the example key to the output
key = Input((1))
# Build the wrapper model
wrapper_model = Model([model.input, key], [model.output, key])
In [ ]:
# Let's grap an arbitrary example from the test data.
example = x_test[7]
# Eventhough it is one example, we need to fit it into a batch (of one)
example = np.expand_dims(example, axis=0)
print(example.shape)
# Let's make an 1D input for our key and choose a unique value --in this case, we use the number 45
example_key = np.asarray([45])
print(example_key.shape)
Let's make the call now. We will get back a tuple. The first value is the prediction (10 element array --one for each digit) and the example key which is passed through (i.e., 45)
In [ ]:
# Do the prediction
prediction, id = wrapper_model.predict([example, example_key])
print(prediction)
print("The Example Key", id)
We did it. You can see the value 45 in the returned id :)