How to Train and Save a CLASSIFICATION 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 to do classification in Keras for Tensorflow Serving(TF Serving).

If you just want to learn how to train a model and save it for TF Serving, Click here for a tutorial on how to do that for prediction, this one is for classification and is a bit more complex.

Getting Started

We will begin by using the simplest yet probably most popular dataset for training a classification model. It is the IRIS dataset.

Attribute Information:

1. sepal length in cm
2. sepal width in cm
3. petal length in cm
4. petal width in cm
classes:
— Iris Setosa
— Iris Versicolour
— Iris Virginica

I have prepared a repo that includes this dataset so you do not have to go digging around. Click here for it, or just clone it.

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

you will find a pre-saved model in it in directory 1, running the code will save another model in directory 2.

To run the code you must set up your environment which is shown on the github. I have set it up where you can use pipenv. Pipenv is kind of like NPM(if you are familiar with node)

Install pipenv

brew install pipenv

Install Dependencies

cd iris_keras_tensorflow_serving
pipenv install

Run Code

pipenv run python index.py

CODE DIVE TIME!!!

import 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

import numpy as np
import pandas as pd
from keras.models import Sequential
from keras.layers import Dense
from keras.utils import np_utils

from sklearn.preprocessing import LabelEncoder

Set tensorflow session

sess = tf.Session()
K.set_session(sess)
K.set_learning_phase(0)  # all new operations will be in test mode from now on

Set variables and randomize things a little

model_version = "2" #Change this to export different model versions, i.e. 2, ..., 7
epoch = 100 ## the higher this number is the more accurate the prediction will be 5000 is a good number to set it at just takes a while to train

seed = 7
np.random.seed(seed)

Extract Data from csv and format it for training

dataframe = pd.read_csv("iris.csv", header=None)
dataset = dataframe.values
X = dataset[:,0:4].astype(float)
Y = dataset[:,4]
# encode class values as integers
encoder = LabelEncoder()
encoder.fit(Y)
encoded_Y = encoder.transform(Y)
# convert integers to dummy variables (i.e. one hot encoded)
one_hot_labels = np_utils.to_categorical(encoded_Y)
labels = []
for label in Y:
    if label not in labels:
        labels.append(label)

Define and Structure Model

model = Sequential()
model.add(Dense(8, input_dim=4, activation="relu"))
model.add(Dense(3, activation="softmax"))

IMPORTANT!

Here is the Part where there was no documentation. This is necessary to be able to save the model for TF serving in order to use the classify api action. We do this to extract the features for the model so we can specify them for the meta graph variables for tf serving.

We will use these to properly generate signature definitions. TF serving gives some documentation here: https://www.tensorflow.org/serving/signature_defs

but not enough for us to figure out how to do it in keras, so here you go:

serialized_tf_example = tf.placeholder(tf.string, name='tf_example')
feature_configs = {'x': tf.FixedLenFeature(shape=[4], dtype=tf.float32),}
tf_example = tf.parse_example(serialized_tf_example, feature_configs)
x = tf.identity(tf_example['x'], name='x')  # use tf.identity() to assign name
y = model(x)

Compile and Train the model

model.compile(loss='categorical_crossentropy', optimizer="adam", metrics=['accuracy'])
model.fit(X, one_hot_labels, epochs=epoch, batch_size=32)

Classification and Prediction signatures for TF serving

values, indices = tf.nn.top_k(y, len(labels))
table = tf.contrib.lookup.index_to_string_table_from_tensor(tf.constant(labels))
prediction_classes = table.lookup(tf.to_int64(indices))

classification_inputs = tf.saved_model.utils.build_tensor_info(serialized_tf_example)
classification_outputs_classes = tf.saved_model.utils.build_tensor_info(prediction_classes)
classification_outputs_scores = tf.saved_model.utils.build_tensor_info(values)

classification_signature = tf.saved_model.signature_def_utils.classification_signature_def(
    examples=serialized_tf_example,
    classes=prediction_classes,
    scores=values
)

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

Check if these signatures are valid

valid_prediction_signature = tf.saved_model.signature_def_utils.is_valid_signature(prediction_signature)
valid_classification_signature = tf.saved_model.signature_def_utils.is_valid_signature(classification_signature)

if(valid_prediction_signature == False):
    raise ValueError("Error: Prediction signature not valid!")

if(valid_classification_signature == False):
    raise ValueError("Error: Classification signature not valid!")

Save the model

builder = saved_model_builder.SavedModelBuilder('./'+model_version)
legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')

# Add the meta_graph and the variables to the builder
builder.add_meta_graph_and_variables(
      sess, [tag_constants.SERVING],
      signature_def_map={
           'predict-iris':
               prediction_signature,
           signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
              classification_signature,
      },
      legacy_init_op=legacy_init_op)
# save the graph
builder.save()

DONE

There you go that is how to properly save a classification model using keras for tf serving.

Best of luck

— — Brian Alois Schardt