Export a custom Keras model to be used for prediction with the Cloud ML Engine

Question

I have difficulties exporting a custom VGG-Net (not exactly the one from Keras), that was trained with Keras, so that it can be used for the Google Cloud Predict API. I am loading my model with Keras.

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

model = load_model(model.h5)

The image that I want to classify was encoded as base64 string. So, I will have to decode it for the prediction task with some code that I found in one of the google examples.

channels = 3
height = 96
width = 96

def decode_and_resize(image_str_tensor):
   """Decodes jpeg string, resizes it and returns a uint8 tensor."""
   image = tf.image.decode_jpeg(image_str_tensor, channels=channels)
   image = tf.expand_dims(image, 0)
   image = tf.image.resize_bilinear(
       image, [height, width], align_corners=False)
   image = tf.squeeze(image, squeeze_dims=[0])
   image = tf.cast(image, dtype=tf.uint8)
   return image

image_str_tensor = tf.placeholder(tf.string, shape=[None])
key_input = tf.placeholder(tf.string, shape=[None]) 
key_output = tf.identity(key_input)

input_tensor = tf.map_fn(
    decode_and_resize, image_str_tensor, back_prop=False, dtype=tf.uint8)
input_tensor = tf.image.convert_image_dtype(image, dtype=tf.float32)

But after this point, I do no longer know how to proceed. How can I now put this input tensor into my model and get out the correct output tensor, so that I am able to define the SignatureDef and then export my graph as a SavedModel?

Any help would be appreciated.

Answer 1

Disclaimer: Although I'm an expert on Cloud ML Engine's prediction service, and fairly knowledgeable on TensorFlow, I do not know Keras very well. I'm simply piecing together information from other places, notably, this sample and this answer. I can only imagine there being better ways to do this and I hope folks will post such. In the meantime, I hope this serves your needs.

This particular answer assumes you've already saved the model. The code loads the model then exports it as a SavedModel.

The basic idea is to start building a "raw" TensorFlow model for the inputs (the input placeholder, the image decoding, resizing, and batching, etc.), and then "connect" it a Keras VGG model by "rebuilding" the VGG model structure and, finally, restoring the saved weights into the newly built model. Then we save this version of the model out as a SavedModel.

The "magic" here is the connection between the raw TF preprocessing and the VGG model. This happens by passing the "output" of the TF preprocessing graph (input_tensor in the code below) as the input_tensor to the Keras VGG graph.input_tensor contains a batch of already decoded and resized images, just like VGG expects.

import keras.backend as K
import tensorflow as tf
from keras.models import load_model, Sequential
from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import tag_constants, signature_constants
from tensorflow.python.saved_model.signature_def_utils_impl import predict_signature_def

MODEL_FILE = 'model.h5'
WEIGHTS_FILE = 'weights.h5'
EXPORT_PATH = 'YOUR/EXPORT/PATH'

channels = 3
height = 96
width = 96

def build_serving_inputs():

  def decode_and_resize(image_str_tensor):
     """Decodes jpeg string, resizes it and returns a uint8 tensor."""
     image = tf.image.decode_jpeg(image_str_tensor, channels=channels)
     image = tf.expand_dims(image, 0)
     image = tf.image.resize_bilinear(
         image, [height, width], align_corners=False)
     image = tf.squeeze(image, squeeze_dims=[0])
     image = tf.cast(image, dtype=tf.uint8)
     return image

  image_str_tensor = tf.placeholder(tf.string, shape=[None])
  key_input = tf.placeholder(tf.string, shape=[None]) 
  key_output = tf.identity(key_input)

  input_tensor = tf.map_fn(
      decode_and_resize, image_str_tensor, back_prop=False, dtype=tf.uint8)
  input_tensor = tf.image.convert_image_dtype(input_tensor, dtype=tf.float32) 

  return image_str_tensor, input_tensor, key_input, key_output

# reset session
K.clear_session()

with tf.Graph().as_default() as g, tf.Session(graph=g) as sess:
  K.set_session(sess)

  image_str_tensor, input_tensor, key_input, key_output = build_serving_inputs()

  # disable loading of learning nodes
  K.set_learning_phase(0)

  # Load model and save out the weights
  model = load_model(MODEL_FILE)
  model.save_weights(WEIGHTS_FILE)

  # Rebuild the VGG16 model with the weights
  new_model = keras.applications.vgg16.VGG16(
    include_top=True, weights=WEIGHTS_FILE, input_tensor=input_tensor,
    input_shape=[width, height, channels], pooling=None)

  # export saved model
  tf.saved_model.simple_save(
      sess,
      EXPORT_PATH,
      inputs={'image_bytes': image_str_tensor, 'key': key_input},
      outputs={'predictions': new_model.outputs[0], 'key': key_output}
  )

Note I don't know if this code quite works yet (haven't tested); I'm worried about how it handles the batch dimension. build_serving_inputs creates a tensor with a batch dimension and passes it to Keras.