How to Train and Save a PREDICTION Model in KERAS for Tensorflow Serving

So many tutorials on Machine Learning… yet some don’t work, and virtually none of them prepare you to save the model for production use.

Tensorflow serving, is currently the best way to productionize AI, hands down. However, because it is so new, most examples are complex and incomplete.

This is especially true when it comes to saving a model in Keras.

THUS, for the sake of simplicity and to fill this gap I have chose to use keras to train a simple model and export it in the proper format for Tensorflow serving(TF Serving). This tutorial is part of a series, so if you don’t know how to implement TF Serving start here!

Enough Talking.. Lets get started!

XOR Prediction Example

For the sake of clarity I wanted to use the simplest example I could think of, which is the XOR logic gate. https://en.wikipedia.org/wiki/XOR_gate

1. Download code

First clone the Repo to get the code

git clone https://github.com/brianalois/xor_keras_tensorflow_serving.git

Make sure your environment is set up on the github page the readme has instructions.

2. Code Dive

Import YOUR dependencies!!

import tensorflow as tf
from keras import backend as K

from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import tag_constants, signature_constants, signature_def_utils_impl

from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation
from keras.optimizers import SGD
import numpy as np

Set Tensorflow session

sess = tf.Session()
K.set_session(sess)
K.set_learning_phase(0)

Set Variables

model version: will change the name of the exported model directory

epoch: number of training cycles, the larger this is the longer it will take. 5,000 is a good number for accuracy, but may take sometime to train. 100 is good for quick testing.

model_version = "2"
epoch = 100

Load up data

This is why I chose the XOR example because the data is super easy to get

X = np.array([
    [0,0],
    [0,1],
    [1,0],
    [1,1]
])
Y = np.array([[0],[1],[1],[0]])

Build the Model

model = Sequential()
model.add(Dense(8, input_dim=2))
model.add(Activation('tanh'))
model.add(Dense(1))
model.add(Activation('sigmoid'))
sgd = SGD(lr=0.1)

Compile and fit the model

model.compile(loss='binary_crossentropy', optimizer=sgd)
model.fit(X, Y, batch_size=1, nb_epoch=epoch)

Get Tensorflow Serving Input and Output variables

x = model.input
y = model.output
prediction_signature = tf.saved_model.signature_def_utils.predict_signature_def({"inputs": x}, {"prediction":y})

Test if the prediction signature for tf serving is valid

valid_prediction_signature = tf.saved_model.signature_def_utils.is_valid_signature(prediction_signature)
if(valid_prediction_signature == False):
    raise ValueError("Error: Prediction signature not valid!")

Build and confiugure model

builder = saved_model_builder.SavedModelBuilder('./'+model_version)
legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
builder.add_meta_graph_and_variables(
      sess, [tag_constants.SERVING],
      signature_def_map={
           signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:prediction_signature,
      },
      legacy_init_op=legacy_init_op)

Save Model

builder.save()

There you go pretty simple, you should now see a folder in the directory with the name of 2 or whatever the model version you put was.

Note:

For those familiar with Tensorflow serving there are prediction, classification, and regression actions. For each one it requires saving the model in a particular way. Prediction is the easiest to understand which is why I started it. However, classification and regression was difficult to figure out how to get working with keras, so if you are in need of that one let me know and I will publish the instructions.

Peace,

— — Brian Alois Schardt