The model after training has to be exported and saved. The weights, biases, and the graph are stored for inference. We will train an MNIST model and store it. Start with defining the constants that are required, using the following code:
work_dir = '/tmp'
model_version = 9
training_iteration = 1000
input_size = 784
no_classes = 10
batch_size = 100
total_batches = 200
The model_version can be an integer to specify which model we want to export for serving. The feature config is stored as a dictionary with placeholder names and their corresponding datatype. The prediction classes and their labels should be mapped. The identity placeholder can be used with the API:
tf_example = tf.parse_example(tf.placeholder(tf.string, name='tf_example'),
{'x': tf.FixedLenFeature(shape=[784], dtype=tf.float32), })
x_input = tf.identity(tf_example['x'], name='x')
A simple classifier can be defined with weights, biases, logits, and an optimizer, using the following code:
y_input = tf.placeholder(tf.float32, shape=[None, no_classes])
weights = tf.Variable(tf.random_normal([input_size, no_classes]))
bias = tf.Variable(tf.random_normal([no_classes]))
logits = tf.matmul(x_input, weights) + bias
softmax_cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels=y_input, logits=logits)
loss_operation = tf.reduce_mean(softmax_cross_entropy)
optimiser = tf.train.GradientDescentOptimizer(0.5).minimize(loss_operation)
Train the model as shown in the following code:
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
for batch_no in range(total_batches):
mnist_batch = mnist.train.next_batch(batch_size)
_, loss_value = session.run([optimiser, loss_operation], feed_dict={
x_input: mnist_batch[0],
y_input: mnist_batch[1]
})
print(loss_value)
Define the prediction signature, and export the model. Save the model to a persistent storage so that it can be used for inference at a later point in time. This exports the data by deserialization and stores it in a format that can be understood by different systems. Multiple graphs with different variables and placeholders can be used for exporting. It also supports signature_defs and assets. The signature_defs have the inputs and outputs specified because input and output will be accessed from the external clients. Assets are non-graph components that will be utilized for the inference, such as vocabulary and so on.
The classification signature uses access to the classification API of TensorFlow. An input is compulsory and there are two optional outputs (prediction classes and prediction probabilities), with at least one being compulsory. The prediction signature offers flexibility with the number of inputs and outputs. Multiple outputs can be defined and explicitly queried from the client side. The signature_def is shown here:
signature_def = (
tf.saved_model.signature_def_utils.build_signature_def(
inputs={'x': tf.saved_model.utils.build_tensor_info(x_input)},
outputs={'y': tf.saved_model.utils.build_tensor_info(y_input)},
method_name="tensorflow/serving/predict"))
Finally, add the metagraph and variables to the builder with the prediction signature:
model_path = os.path.join(work_dir, str(model_version))
saved_model_builder = tf.saved_model.builder.SavedModelBuilder(model_path)
saved_model_builder.add_meta_graph_and_variables(
session, [tf.saved_model.tag_constants.SERVING],
signature_def_map={
'prediction': signature_def
},
legacy_init_op=tf.group(tf.tables_initializer(), name='legacy_init_op'))
saved_model_builder.save()
The builder is saved and ready to be consumed by the server. The shown example is applicable to any model and can be used for exporting. In the next section, we will serve and query the exported model.