<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Prepare and Deploy a TensorFlow Model to AI Platform for Online Serving
This Notebook demonstrates how to prepare a TensorFlow 2.x model and deploy it for serving with AI Platform Prediction. This example uses the pretrained ResNet V2 101 image classification model from TensorFlow Hub (TF Hub).
The Notebook covers the following steps
Step1:  Configure GCP environment settings
Step2:  Create a local workspace
Step3:  1. Loading and Running the ResNet Module
1.1. Download and instantiate the model
Step4:  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
Step5:  1.3. Preprocess the testing images
The images need to be preprocessed to conform to the format expected by the ResNet101 model.
Step6:  2.4. Run inference
Step7:  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.
Step8:  We will display the 5 highest ranked labels for each image
Step9:  2. Create Serving Signatures
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.
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.
Step16:  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
Step17:  2.2. Test the custom serving module
Step18:  3. Save the custom serving module as SavedModel
Step19:  3.1. Inspect the SavedModel
Step20:  3.2. Test loading and executing the SavedModel
Step21:  4. Deploy the SavedModel to AI Platform Prediction
4.1. Copy the SavedModel to GCS
Step22:  4.2 Create a model in AI Platform Prediction
Step23:  4.3 Create a model version
Step24:  5. Validate the Deployed Model Version to AI Platform Prediction
<ASSISTANT_TASK:>
Python 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
Explanation: Prepare and Deploy a TensorFlow Model to AI Platform for Online Serving
This Notebook demonstrates how to prepare a TensorFlow 2.x model and deploy it for serving with AI Platform Prediction. This example uses the pretrained ResNet V2 101 image classification model from TensorFlow Hub (TF Hub).
The Notebook covers the following steps:
1. Downloading and running the ResNet module from TF Hub
2. Creating serving signatures for the module
3. Exporting the model as a SavedModel
4. Deploying the SavedModel to AI Platform Prediction
5. Validating the deployed model
Setup
This Notebook was tested on AI Platform Notebooks using the standard TF 2.2 image.
Import libraries
End of explanation
PROJECT_ID = '[your-google-project-id]' # Set your project Id
BUCKET = '[your-bucket-name]' # Set your bucket name Id
REGION = '[your-region]'  # Set your region for deploying the model
MODEL_NAME = 'resnet_classifier'
MODEL_VERSION = 'v1'
GCS_MODEL_LOCATION = 'gs://{}/models/{}/{}'.format(BUCKET, MODEL_NAME, MODEL_VERSION)
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'
IMAGES_FOLDER = 'test_images'
!gcloud config set project $PROJECT_ID
Explanation: Configure GCP environment settings
End of explanation
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)
Explanation: Create a local workspace
End of explanation
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)
Explanation: 1. Loading and Running the ResNet Module
1.1. Download and instantiate the model
End of explanation
image_list = [tf.io.read_file(os.path.join(IMAGES_FOLDER, image_path))
         for image_path in os.listdir(IMAGES_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())
Explanation: 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 with class 0 for "background", followed by 1000 actual ImageNet classes.
We will now test the model on a couple of JPEG images. 
1.2. Display sample images
End of explanation
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)
Explanation: 1.3. Preprocess the testing images
The images need to be preprocessed to conform to the format expected by the ResNet101 model.
End of explanation
predictions = model(preprocessed_images)
predictions
Explanation: 2.4. Run inference
End of explanation
labels_path = tf.keras.utils.get_file(
    'ImageNetLabels.txt',
    IMAGENET_LABELS_URL)
imagenet_labels = np.array(open(labels_path).read().splitlines())
Explanation: 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.
End of explanation
for prediction in list(predictions):
    decoded = imagenet_labels[np.argsort(prediction.numpy())[::-1][:5]]
    print(list(decoded))
Explanation: We will display the 5 highest ranked labels for each image
End of explanation
list(model.signatures)
Explanation: 2. Create Serving Signatures
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.
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.
End of explanation
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)
Explanation: 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: 
1. The default signature - This will expose the default predict method of the ResNet101 model.
2. Prep/post-processing signature - 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. 
The signatures are created 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.
2.1. Define the custom serving module
End of explanation
predictions = serving_module.predict_labels(raw_images)
predictions
Explanation: 2.2. Test the custom serving module
End of explanation
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)
Explanation: 3. Save the custom serving module as SavedModel
End of explanation
!saved_model_cli show --dir {model_path} --tag_set serve --all
Explanation: 3.1. Inspect the SavedModel
End of explanation
model = tf.keras.models.load_model(model_path)
model.predict_labels(raw_images)
Explanation: 3.2. Test loading and executing the SavedModel
End of explanation
!gsutil cp -r {model_path} {GCS_MODEL_LOCATION}
!gsutil ls {GCS_MODEL_LOCATION}
Explanation: 4. Deploy the SavedModel to AI Platform Prediction
4.1. Copy the SavedModel to GCS
End of explanation
!gcloud ai-platform models create {MODEL_NAME} \
  --project {PROJECT_ID} \
  --regions {REGION}
!gcloud ai-platform models list --project {PROJECT_ID} 
Explanation: 4.2 Create a model in AI Platform Prediction
End of explanation
MACHINE_TYPE='n1-standard-8'
ACCELERATOR='count=1,type=nvidia-tesla-p4'
!gcloud beta ai-platform versions create {MODEL_VERSION} \
  --model={MODEL_NAME} \
  --origin={GCS_MODEL_LOCATION} \
  --runtime-version=2.1 \
  --framework=TENSORFLOW \
  --python-version=3.7 \
  --machine-type={MACHINE_TYPE} \
  --accelerator={ACCELERATOR} \
  --project={PROJECT_ID}
!gcloud ai-platform versions list --model={MODEL_NAME} --project={PROJECT_ID}
Explanation: 4.3 Create a model version
End of explanation
import googleapiclient.discovery
service = googleapiclient.discovery.build('ml', 'v1')
name = 'projects/{}/models/{}/versions/{}'.format(PROJECT_ID, MODEL_NAME, MODEL_VERSION)
print("Service name: {}".format(name))
def caip_predict(instances, signature_name='serving_default'):
    
    request_body={
        'signature_name': signature_name,
        'instances': instances}
    
    response = service.projects().predict(
        name=name,
        body=request_body
    
    ).execute()
    if 'error' in response:
        raise RuntimeError(response['error'])
    outputs = response['predictions']
    return outputs
signature_name = 'serving_preprocess'
encoded_images = [{'b64': base64.b64encode(image.numpy()).decode('utf-8')} 
                  for image in image_list]  
caip_predict(encoded_images, signature_name=signature_name)
Explanation: 5. Validate the Deployed Model Version to AI Platform Prediction
End of explanation
<END_TASK>