This is a demo that demonstrates end to end workflow from training to serving of MNIST model with Tensorflow and Kubernetes



In [ ]:

    
import os
import sys

# This is a placeholder for a Google-internal import.

import tensorflow as tf

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

# Base path for exporting model
# Set it to /home/jovyan/tf-serving-base/models to easily create a container and deploy
export_path_base = "/home/jovyan/tf-serving-base/models"
# Model Version
model_version = 1

Let's look at the devices available locally



In [ ]:

    
from tensorflow.python.client import device_lib
local_device_protos = device_lib.list_local_devices()
[print(x.name) for x in local_device_protos if x.device_type == 'GPU']

Let's train a model based on mnist data



In [ ]:

    
# Training iterations
training_iteration = 1000

print('Training model...')
sess = tf.InteractiveSession()
serialized_tf_example = tf.placeholder(tf.string, name='tf_example')
feature_configs = {'x': tf.FixedLenFeature(shape=[784], 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_ = tf.placeholder('float', shape=[None, 10])
w = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
sess.run(tf.global_variables_initializer())
y = tf.nn.softmax(tf.matmul(x, w) + b, name='y')
cross_entropy = -tf.reduce_sum(y_ * tf.log(y))
train_step = tf.train.GradientDescentOptimizer(0.01).minimize(cross_entropy)
values, indices = tf.nn.top_k(y, 10)
table = tf.contrib.lookup.index_to_string_table_from_tensor(
    tf.constant([str(i) for i in range(10)]))
prediction_classes = table.lookup(tf.to_int64(indices))
for _ in range(training_iteration):
    batch = mnist.train.next_batch(50)
    train_step.run(feed_dict={x: batch[0], y_: batch[1]})
correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, 'float'))
print('training accuracy %g' % sess.run(
    accuracy, feed_dict={x: mnist.test.images,
                         y_: mnist.test.labels}))
print('Done training!')

Let's export the model to a pre-defined directory.



In [ ]:

    
export_path = os.path.join(
    tf.compat.as_bytes(export_path_base),
    tf.compat.as_bytes(str(model_version)))
print('Exporting trained model to', export_path)
builder = tf.saved_model.builder.SavedModelBuilder(export_path)

# Build the signature_def_map.
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.build_signature_def(
          inputs={
              tf.saved_model.signature_constants.CLASSIFY_INPUTS:
                  classification_inputs
          },
          outputs={
              tf.saved_model.signature_constants.CLASSIFY_OUTPUT_CLASSES:
                  classification_outputs_classes,
              tf.saved_model.signature_constants.CLASSIFY_OUTPUT_SCORES:
                  classification_outputs_scores
          },
          method_name=tf.saved_model.signature_constants.CLASSIFY_METHOD_NAME))

tensor_info_x = tf.saved_model.utils.build_tensor_info(x)
tensor_info_y = tf.saved_model.utils.build_tensor_info(y)

prediction_signature = (
      tf.saved_model.signature_def_utils.build_signature_def(
          inputs={'images': tensor_info_x},
          outputs={'scores': tensor_info_y},
          method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))

legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
builder.add_meta_graph_and_variables(
      sess, [tf.saved_model.tag_constants.SERVING],
      signature_def_map={
          'predict_images':
              prediction_signature,
          tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
              classification_signature,
      },
      legacy_init_op=legacy_init_op)

builder.save()

print('Done exporting!')

Let's build a container with the model just generated.

Open Terminal and run "cd tf-serving-base && ./build_and_export.sh 'gcr.io project name'"

Now lets start a Kubernetes Deployment that runs Tensorflow Serving with our model



In [ ]:

    
from os import path

import yaml
import time

from kubernetes import client, config

config.load_incluster_config()

with open("/home/jovyan/tf-serving-base/deployment.yaml") as f:
    dep = yaml.load(f)
    # Specify the newly built image.
    dep['spec']['template']['spec']['containers'][0]['image'] = "gcr.io/vishnuk-cloud/tf-model-server:latest"
    k8s_beta = client.ExtensionsV1beta1Api()
    try:
        resp = k8s_beta.delete_namespaced_deployment(body={}, namespace="jupyterhub", name=dep['metadata']['name'])
        time.sleep(10)
    except:
        pass
    resp = k8s_beta.create_namespaced_deployment(body=dep, namespace="jupyterhub")
    print("Deployment created. status='%s'" % str(resp.status))

with open("/home/jovyan/tf-serving-base/service.yaml") as f:
    dep = yaml.load(f)
    v1 = client.CoreV1Api()
    v1.delete_namespaced_service(namespace="jupyterhub", name=dep['metadata']['name'])
    resp = v1.create_namespaced_service(body=dep, namespace="jupyterhub")
    print("Service created. status='%s'" % str(resp.status))



In [ ]: