<jupyter_start><jupyter_text>Preparing a TensorFlow 2 model for serving with TF ServingThis notebook demonstrates how to prepare a TensorFlow 2.x model for serving with AI Platform Prediction Custom Containers (Beta) and [TF Serving](https://www.tensorflow.org/tfx/guide/serving). The inputs and outputs of the model as used during model training may not be optimal for serving. For example, in a typical training pipeline, feature engineering is performed as a separate step preceding model training and hyperparameter tuning. When serving the model, it may be more optimal to embed the feature engineering logic into the serving interface rather than require a client application to preprocess data.This notebook uses the pretrained [ResNet V2 101](https://tfhub.dev/google/imagenet/resnet_v2_101/classification/4) image classification model from [TF Hub](https://tfhub.dev/). The demonstrated patterns and practices can be easily generalized to other types of TensorFlow 2 models.<jupyter_code>import base64
import os
import json
import requests
import time
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import matplotlib.pyplot as plt

from typing import List, Optional, Text, Tuple<jupyter_output><empty_output><jupyter_text>Setting up the environmentThis notebook was tested on **AI Platform Notebooks** using the standard TF 2.2 image.   Set the model store pathUpdate the `GCS_MODEL_STORE` constant with the name of the GCS path to store the model created later in this notebook.<jupyter_code>GCS_MODEL_STORE = 'gs://mlops-dev-workspace/models'<jupyter_output><empty_output><jupyter_text>Set the urls to TF Hub assets<jupyter_code>THUB_MODEL_HANDLE = 'https://tfhub.dev/google/imagenet/resnet_v2_101/classification/4'
IMAGENET_LABELS_URL = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'<jupyter_output><empty_output><jupyter_text>Create a local workspace<jupyter_code>LOCAL_WORKSPACE = '/tmp/workspace'
if tf.io.gfile.exists(LOCAL_WORKSPACE):
  print("Removing previous workspace artifacts...")
  tf.io.gfile.rmtree(LOCAL_WORKSPACE)

print("Creating a new workspace...")
tf.io.gfile.makedirs(LOCAL_WORKSPACE)<jupyter_output>Removing previous workspace artifacts...
Creating a new workspace...<jupyter_text>Loading and running the ResNet v2 101 model  Download and instantiate the model<jupyter_code>os.environ["TFHUB_DOWNLOAD_PROGRESS"] = 'True'

local_savedmodel_path = hub.resolve(THUB_MODEL_HANDLE)

print(local_savedmodel_path)
!ls -la {local_savedmodel_path}
model = hub.load(THUB_MODEL_HANDLE)<jupyter_output><empty_output><jupyter_text>The expected input to most TF Hub TF2 image classification models, including ResNet 101, is a rank 4 tensor conforming to the following tensor specification: `tf.TensorSpec([None, height, width, 3], tf.float32)`. For the ResNet 101 model, the expected image size is `height x width = 224 x 224`. The color values for all channels are expected to be normalized to the [0, 1] range. The output of the model is a batch of logits vectors. The indices into the logits are the `num_classes = 1001` classes  from the ImageNet dataset. The mapping from indices to class labels can be found in the [labels file](download.tensorflow.org/data/ImageNetLabels.txt) with class 0 for "background", followed by 1000 actual ImageNet classes.      Test the modelWe will now test the model on a couple of JPEG images.   Display sample images<jupyter_code>image_folder = 'test_images'
image_list = [tf.io.read_file(os.path.join(image_folder, image_path))
         for image_path in os.listdir(image_folder)]

ncolumns = len(image_list) if len(image_list) < 4 else 4
nrows = int(len(image_list) // ncolumns)
fig, axes = plt.subplots(nrows=nrows, ncols=ncolumns, figsize=(10,10))
for axis, image in zip(axes.flat[0:], image_list):
    decoded_image = tf.image.decode_image(image)
    axis.set_title(decoded_image.shape)
    axis.imshow(decoded_image.numpy())<jupyter_output><empty_output><jupyter_text>Preprocess the testing imagesThe images need to be preprocessed to conform to the format expected by the ResNet101 model.<jupyter_code>def _decode_and_scale(image, size):
    image = tf.image.decode_image(image, expand_animations=False)
        
    image_height = image.shape[0]
    image_width = image.shape[1]
    crop_size = tf.minimum(image_height, image_width)
    offset_height = ((image_height - crop_size) + 1) // 2
    offset_width = ((image_width - crop_size) + 1) // 2
        
    image = tf.image.crop_to_bounding_box(image, offset_height, offset_width, crop_size, crop_size)
    image = tf.cast(tf.image.resize(image, [size, size]), tf.uint8)
    
    return image
size = 224

raw_images = tf.stack(image_list)
preprocessed_images = tf.map_fn(lambda x: _decode_and_scale(x, size), raw_images, dtype=tf.uint8)
preprocessed_images = tf.image.convert_image_dtype(preprocessed_images, tf.float32)
print(preprocessed_images.shape)<jupyter_output>(2, 224, 224, 3)<jupyter_text>Run inference<jupyter_code>predictions = model(preprocessed_images)
predictions<jupyter_output><empty_output><jupyter_text>The model returns a batch of arrays with logits. This is not a very user friendly output so we will convert it to the list of ImageNet class labels.  Download ImageNet labels<jupyter_code>labels_path = tf.keras.utils.get_file(
    'ImageNetLabels.txt',
    IMAGENET_LABELS_URL)
imagenet_labels = np.array(open(labels_path).read().splitlines())<jupyter_output><empty_output><jupyter_text>Map the logits to class labelsWe will display the 5 highest ranked labels for each image<jupyter_code>for prediction in list(predictions):
    decoded = imagenet_labels[np.argsort(prediction.numpy())[::-1][:5]]
    print(list(decoded))<jupyter_output>['military uniform', 'suit', 'Windsor tie', 'pickelhaube', 'bow tie']
['Egyptian cat', 'tiger cat', 'tabby', 'lynx', 'Siamese cat']<jupyter_text>Preparing the model for TF Serving The ResNet V2 101 model from TF HUB is optimized for recomposition and fine tuning. Since there are no serving signatures in the model's metadata, it cannot be served with TF Serving as is.<jupyter_code>list(model.signatures)<jupyter_output><empty_output><jupyter_text>To make it servable, we need to add a serving signature(s) describing the inference method(s) of the model. We will add two signatures. The default signature will expose the default predict method of the ResNet101 model.Since the expected inputs to this interface require a relatively complex image preprocessing to be performed by a client, we will also expose an alternative signature that embeds the preprocessing and postprocessing logic and accepts raw unprocessed images and returns the list of ranked class labels and associated label probabilities. This is achieved by defining a custom module class derived from the `tf.Module` base class that encapsulates our ResNet model and extends it with a method implementing the image preprocessing and output postprocessing logic. The default method of the custom module is mapped to the default method of the base ResNet module to maintain the analogous interface. The custom module will be exported as `SavedModel` that includes the original model, the preprocessing logic, and two serving signatures.This technique can be generalized to other scenarios where you need to extend a TensorFlow model and you have access to the serialized `SavedModel` but you don't have access to the Python code implementing the model.  Define the custom serving module<jupyter_code>LABELS_KEY = 'labels'
PROBABILITIES_KEY = 'probabilities'
NUM_LABELS = 5

class ServingModule(tf.Module):
    """
    A custom tf.Module that adds image preprocessing and output post processing to
    a base TF 2 image classification model from TF Hub. 
    """

    def __init__(self, base_model, input_size, output_labels):
        super(ServingModule, self).__init__()
        self._model = base_model
        self._input_size = input_size
        self._output_labels = tf.constant(output_labels, dtype=tf.string)
        

    def _decode_and_scale(self, raw_image):
        """
        Decodes, crops, and resizes a single raw image.
        """
        
        image = tf.image.decode_image(raw_image, dtype=tf.dtypes.uint8, expand_animations=False)
        image_shape = tf.shape(image)
        image_height = image_shape[0]
        image_width = image_shape[1]
        crop_size = tf.minimum(image_height, image_width)
        offset_height = ((image_height - crop_size) + 1) // 2
        offset_width = ((image_width - crop_size) + 1) // 2
        
        image = tf.image.crop_to_bounding_box(image, offset_height, offset_width, crop_size, crop_size)
        image = tf.image.resize(image, [self._input_size, self._input_size])
        image = tf.cast(image, tf.uint8)
    
        return image
    
    def _preprocess(self, raw_inputs):
        """
        Preprocesses raw inputs as sent by the client.
        """
        
        # A mitigation for https://github.com/tensorflow/tensorflow/issues/28007
        with tf.device('/cpu:0'):
            images = tf.map_fn(self._decode_and_scale, raw_inputs, dtype=tf.uint8)
        images = tf.image.convert_image_dtype(images, tf.float32)
        
        return images
        
    def _postprocess(self, model_outputs):
        """
        Postprocesses outputs returned by the base model.
        """
        
        probabilities = tf.nn.softmax(model_outputs)
        indices = tf.argsort(probabilities, axis=1, direction='DESCENDING')
        
        return {
            LABELS_KEY: tf.gather(self._output_labels, indices, axis=-1)[:,:NUM_LABELS],
            PROBABILITIES_KEY: tf.sort(probabilities, direction='DESCENDING')[:,:NUM_LABELS]
        }
        

    @tf.function(input_signature=[tf.TensorSpec([None, 224, 224, 3], tf.float32)])
    def __call__(self, x):
        """
        A pass-through to the base model.
        """
        
        return self._model(x)

    @tf.function(input_signature=[tf.TensorSpec([None], tf.string)])
    def predict_labels(self, raw_images):
        """
        Preprocesses inputs, calls the base model 
        and postprocesses outputs from the base model.
        """
        
        # Call the preprocessing handler
        images = self._preprocess(raw_images)
        
        # Call the base model
        logits = self._model(images)
        
        # Call the postprocessing handler
        outputs = self._postprocess(logits)
        
        return outputs
        
    
serving_module = ServingModule(model, 224, imagenet_labels)<jupyter_output><empty_output><jupyter_text>Test the custom serving module<jupyter_code>predictions = serving_module.predict_labels(raw_images)
predictions<jupyter_output><empty_output><jupyter_text>Save the custom serving module as `SavedModel`<jupyter_code>model_name = 'resnet_serving'
model_version = '1'
model_path = os.path.join(LOCAL_WORKSPACE, model_name, model_version)

default_signature = serving_module.__call__.get_concrete_function()
preprocess_signature = serving_module.predict_labels.get_concrete_function()
signatures = {
    'serving_default': default_signature,
    'serving_preprocess': preprocess_signature
}

tf.saved_model.save(serving_module, model_path, signatures=signatures)<jupyter_output>INFO:tensorflow:Assets written to: /tmp/workspace/resnet_serving/1/assets<jupyter_text>Inspect the `SavedModel`<jupyter_code>!saved_model_cli show --dir {model_path} --tag_set serve --all<jupyter_output>MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['__saved_model_init_op']:
  The given SavedModel SignatureDef contains the following input(s):
  The given SavedModel SignatureDef contains the following output(s):
    outputs['__saved_model_init_op'] tensor_info:
        dtype: DT_INVALID
        shape: unknown_rank
        name: NoOp
  Method name is: 

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['x'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 224, 224, 3)
        name: serving_default_x:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output_0'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1001)
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict

signature_def['serving_preprocess']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['raw_images'] tensor[...]<jupyter_text>Test the custom serving module with TF Serving We will now run the exported custom serving module using the TF Serving docker image. Start the TF Serving server:To start the *TF Serving* server and expose the *TF Serving* REST API port (8501):1. Open a JupyterLab terminal.2. Pull the latest TF Serving docker image.```docker pull tensorflow/serving```4. Set the environment variables for the model's name and the model's path. Use the value printed by the next cell for the model's path.```export MODEL_PATH=[YOUR_MODEL_PATH]export MODEL_NAME=resnet```<jupyter_code>print(model_path[:-2])<jupyter_output>/tmp/workspace/resnet_serving<jupyter_text>3. Start the server and register your custom module for serving```docker run -it --rm -p 8501:8501 \-v $MODEL_PATH:/models/$MODEL_NAME \-e MODEL_NAME=$MODEL_NAME \tensorflow/serving ```  Run inference We will now run inference by invoking the TF Serving `Predict` API.Refer to the [TF Serving REST API Reference](https://www.tensorflow.org/tfx/serving/api_rest) for more information about the API format.  Invoke the model using the `serving_preprocess` signature Load sample images<jupyter_code>image_folder = 'test_images'
image_list = [tf.io.read_file(os.path.join(image_folder, image_path))
         for image_path in os.listdir(image_folder)]<jupyter_output><empty_output><jupyter_text>Encode the images using base 64 encoding.<jupyter_code>raw_images_as_list = [{'b64': base64.b64encode(image.numpy()).decode('utf-8')} for image in image_list]<jupyter_output><empty_output><jupyter_text>Prepare the request body and headers.<jupyter_code>request_body = {
    'signature_name': 'serving_preprocess',
    'instances': raw_images_as_list
}

headers = {"content-type": "application/json"}<jupyter_output><empty_output><jupyter_text>Invoke the `Predict` endpoint and display results.<jupyter_code>uri = 'http://localhost:8501/v1/models/resnet:predict'

response = requests.post(uri, data=json.dumps(request_body))

response.json()<jupyter_output><empty_output><jupyter_text>Invoke the model using the `serving_default` signature Preprocess the raw images.<jupyter_code>size = 224

raw_images = tf.stack(image_list)
preprocessed_images = tf.map_fn(lambda x: _decode_and_scale(x, size), raw_images, dtype=tf.uint8)
preprocessed_images = tf.image.convert_image_dtype(preprocessed_images, tf.float32)
print(preprocessed_images.shape)<jupyter_output>(2, 224, 224, 3)<jupyter_text>Prepare the request body and headers.<jupyter_code>images_as_list = preprocessed_images.numpy().tolist()

request_body = {
    'signature_name': 'serving_default',
    'instances': images_as_list
}

headers = {"content-type": "application/json"}<jupyter_output><empty_output><jupyter_text>Invoke the `Predict` endpoint and display results.<jupyter_code>uri = 'http://localhost:8501/v1/models/resnet:predict'

response = requests.post(uri, data=json.dumps(request_body))

print(str(response.json())[0:200], '...')<jupyter_output>{'predictions': [[0.28163141, 0.614411712, -0.00329366326, 0.873328269, -0.783175826, -0.431026816, -0.0673072636, -0.836482167, -0.432366759, 0.229343504, -1.22413266, -0.388189197, 0.706549644, -0.0 ...<jupyter_text>Stop the TF Serving container1. Open a new JupyterLab terminal.2. Get the ID of your TF Serving container```docker ps```3. Terminate your TF Serving container```docker kill [YOUR_CONTAINER_ID]```  Copy the custom serving module to GCS<jupyter_code>local_path = os.path.join(LOCAL_WORKSPACE, model_name)

!gsutil cp -r {local_path} {GCS_MODEL_STORE}
!gsutil ls {GCS_MODEL_STORE}<jupyter_output>gs://mlops-dev-workspace/models/covertype_classifier/
gs://mlops-dev-workspace/models/resnet_serving/
{'fraction_non_alphanumeric': 0.09176921649867045, 'fraction_numerical': 0.020716096716344074, 'mean_word_length': 4.755160142348754, 'pattern_counts': {'":': 2, '<': 89, '<?xml version=': 0, '>': 88, 'https://': 7, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '24', 'dir': 'jupyter-structured-clean-dedup', 'id': '273613', 'license': 'Apache-2.0', 'n_tokens_mistral': 5220, 'n_tokens_neox': 4913, 'n_words': 1422, 'path': 'model_serving/tfserving-caip/01-prepare-for-serving.ipynb', 'repo_name': 'jarokaz/caipp-serving', 'size': '16171'}