Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save wolf1986/165ebff71b2cbb0d27609faf77238832 to your computer and use it in GitHub Desktop.
Save wolf1986/165ebff71b2cbb0d27609faf77238832 to your computer and use it in GitHub Desktop.

Tensorflow Serving Tutorial - 02 - Using a Custom Model

In this tutorial:

  • Assuming you have an environment that can works with the official inception-v3
  • Adapt an existing model to work with TF Serving
  • Query the server with images and retrieve classification results

Save a model after training

In this example we'll suppose that we want our TF Model to accept base64 strings of image files, i.e. prepared by python:

  with open(PATH_IMAGE, 'rb') as f:
    data = f.read()
    data_base64 = base64.urlsafe_b64encode(data)

When preparing our model, it's possible to preprocess the image as part of the network itself (normalize mean, resize, expand dimensions etc..).

It's even possible to decode the image format from raw bytes, or even from base64:

    base64_decoder = tf.decode_base64(base64_tensor, name='base64_decoder')
    image_reader = tf.image.decode_image(base64_decoder, name='image_reader')
    float_caster = tf.cast(image_reader, tf.float32)
    dims_expander = tf.expand_dims(float_caster, 0)
    # Resize the image, normalize, etc ...

Export model for TF Model Server

To understand the given example of exporting the pre-trained inception model, we'll break down the important part:

        # Obtain a tf.Session
        sess = ... 
        # Obtain pointers to input tensors and output tensors ...
        input_tensor = ...
        output_tensor = ...

        # Extract tensors info for builder
        input_tensor_info = tf.saved_model.utils.build_tensor_info(input_tensor)
        output_tensor_info = tf.saved_model.utils.build_tensor_info(output_tensor)

TF Model server & TF Serving Client need to be aware of the input & output that they expect from each other. In this example, our model performs classification of images. So we'll define a signature definition for the client & the server. Input will be labeled 'images' and the output will be labeled 'classification', and they will be mapped to our input & output tensors respectively.

        ...
        
        builder = tf.saved_model.builder.SavedModelBuilder(output_path)
        legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
  
        prediction_signature = (
            tf.saved_model.signature_def_utils.build_signature_def(
                inputs={'images': input_tensor_info},
                outputs={
                    'classification': output_tensor_info
                },
                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME
            )
        )

        builder.add_meta_graph_and_variables(
            sess, [tf.saved_model.tag_constants.SERVING],
            signature_def_map={
                'predict_images':
                    prediction_signature,
            },
            legacy_init_op=legacy_init_op
        )

        builder.save()

Adapting the client

Testing the Model Server

  • Launch the server with the custom exported model - Same as with the official Inception-v3 model.
  • Launch the client adapter for our custom model
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment