<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 TF-TRT Dynamic shapes example
Step1:  Define the model
Step2:  Prepare data
Instructions to download can be found in the link above.
Define input paths
Step5:  Helper class to load the images and convert to a fixed size
Step6:  Define train and test sets.
Step7:  Train the model
Step8:  Save the trained model in savedmodel format
Step9:  Visualize predictions
Step10:  Helper functions
Step11:  Run inference
Step12:  Now prepare input data with different resolution and run inference using that
Step13:  We have seen so far that a fully convolutional TF model can be used to infer images with different input sizes.
Convert the model with TF-TRT
In dynamic shape mode we will have a single engine that can handle the input shapes seen during build mode. We select profile strategy 'Range', and we need to define the range of input shapes.
Step14:  Load converted model and run inference
Step15:  Predict with different image size.
<ASSISTANT_TASK:>
Python Code:
import os
os.environ["TF_CPP_VMODULE"]="trt_engine_utils=2,segment=2" # verbose output
os.environ["TF_TRT_OP_DENYLIST"] = "Shape"                  # workaround for shape output memory error
!pip install pillow matplotlib
import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt
import numpy as np
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.framework import convert_to_constants
import time
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.preprocessing.image import load_img
from IPython.display import Image, display
from tensorflow.keras.preprocessing.image import load_img
import PIL
from PIL import ImageOps
import matplotlib.pyplot as plt
Explanation: TF-TRT Dynamic shapes example: UNet
The goal of this notebook is to demonstrate TF-TRT support for converting a network that has dynamic image size. While earlier TRT versions allowed only the batch size to be unknown, with dynamic shape mode other dimensions can be dynamic.
We take a fully convolutional network, UNet as an example.
The network and many of the code samples in this notebook are from https://keras.io/examples/vision/oxford_pets_image_segmentation
Requirements
This notebook requires at least TF 2.5 and TRT 7.1.3.
End of explanation
def get_model(img_size, num_classes):
    inputs = keras.Input(shape=img_size + (3,))
    ### [First half of the network: downsampling inputs] ###
    # Entry block
    x = layers.Conv2D(32, 3, strides=2, padding="same")(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    previous_block_activation = x  # Set aside residual
    # Blocks 1, 2, 3 are identical apart from the feature depth.
    for filters in [64, 128, 256]: #[16, 32, 64]:
        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)
        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)
        x = layers.MaxPooling2D(3, strides=2, padding="same")(x)
        # Project residual
        residual = layers.Conv2D(filters, 1, strides=2, padding="same")(
            previous_block_activation
        )
        x = layers.add([x, residual])  # Add back residual
        previous_block_activation = x  # Set aside next residual
    ### [Second half of the network: upsampling inputs] ###
    for filters in [256, 128, 64, 32]: #[64, 32, 16, 8]:
        x = layers.Activation("relu")(x)
        x = layers.Conv2DTranspose(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)
        x = layers.Activation("relu")(x)
        x = layers.Conv2DTranspose(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)
        x = layers.UpSampling2D(2)(x)
        # Project residual
        residual = layers.UpSampling2D(2)(previous_block_activation)
        residual = layers.Conv2D(filters, 1, padding="same")(residual)
        x = layers.add([x, residual])  # Add back residual
        previous_block_activation = x  # Set aside next residual
    # Add a per-pixel classification layer
    outputs = layers.Conv2D(num_classes, 3, activation="softmax", padding="same")(x)
    # Define the model
    model = keras.Model(inputs, outputs)
    return model
# Free up RAM in case the model definition cells were run multiple times
keras.backend.clear_session()
# Build model with uknwown image size
img_size = (None, None)
num_classes = 3
model = get_model(img_size, num_classes)
model.summary()
Explanation: Define the model
End of explanation
input_dir = "unet_pets/images/"
target_dir = "unet_pets/annotations/trimaps/"
input_img_paths = sorted(
    [
        os.path.join(input_dir, fname)
        for fname in os.listdir(input_dir)
        if fname.endswith(".jpg")
    ]
)
target_img_paths = sorted(
    [
        os.path.join(target_dir, fname)
        for fname in os.listdir(target_dir)
        if fname.endswith(".png") and not fname.startswith(".")
    ]
)
print("Number of samples:", len(input_img_paths))
# Just print a few image paths
for input_path, target_path in zip(input_img_paths[:10], target_img_paths[:10]):
    print(input_path, "|", target_path)
Explanation: Prepare data
Instructions to download can be found in the link above.
Define input paths
End of explanation
class OxfordPets(keras.utils.Sequence):
    Helper to iterate over the data (as Numpy arrays).
    def __init__(self, batch_size, img_size, input_img_paths, target_img_paths):
        self.batch_size = batch_size
        self.img_size = img_size
        self.input_img_paths = input_img_paths
        self.target_img_paths = target_img_paths
    def __len__(self):
        return len(self.target_img_paths) // self.batch_size
    def __getitem__(self, idx):
        Returns tuple (input, target) correspond to batch #idx.
        i = idx * self.batch_size
        batch_input_img_paths = self.input_img_paths[i : i + self.batch_size]
        batch_target_img_paths = self.target_img_paths[i : i + self.batch_size]
        x = np.zeros((self.batch_size,) + self.img_size + (3,), dtype="float32")
        for j, path in enumerate(batch_input_img_paths):
            img = load_img(path, target_size=self.img_size)
            x[j] = img
        y = np.zeros((self.batch_size,) + self.img_size + (1,), dtype="uint8")
        for j, path in enumerate(batch_target_img_paths):
            img = load_img(path, target_size=self.img_size, color_mode="grayscale")
            y[j] = np.expand_dims(img, 2)
            # Ground truth labels are 1, 2, 3. Subtract one to make them 0, 1, 2:
            y[j] -= 1
        return x, y
Explanation: Helper class to load the images and convert to a fixed size
End of explanation
import random
def train_test_split(input_img_paths, target_img_paths):
    # Split our img paths into a training and a validation set
    val_samples = 1000
    random.Random(1337).shuffle(input_img_paths)
    random.Random(1337).shuffle(target_img_paths)
    train_input_img_paths = input_img_paths[:-val_samples]
    train_target_img_paths = target_img_paths[:-val_samples]
    val_input_img_paths = input_img_paths[-val_samples:]
    val_target_img_paths = target_img_paths[-val_samples:]
    return train_input_img_paths, train_target_img_paths, val_input_img_paths, val_target_img_paths
train_inp_paths, train_target_paths, val_inp_paths, val_target_paths = \
    train_test_split(input_img_paths, target_img_paths)
# Instantiate data Sequences for each split
img_size = (160, 160)
batch_size = 4
train_gen = OxfordPets(batch_size, img_size, train_inp_paths, train_target_paths)
val_gen = OxfordPets(batch_size, img_size, val_inp_paths, val_target_paths)
Explanation: Define train and test sets.
End of explanation
# Configure the model for training.
# We use the "sparse" version of categorical_crossentropy
# because our target data is integers.
model.compile(optimizer="rmsprop", loss="sparse_categorical_crossentropy")
callbacks = [
    keras.callbacks.ModelCheckpoint("oxford_segmentation.h5", save_best_only=True)
]
# Train the model, doing validation at the end of each epoch.
epochs = 15
model.fit(train_gen, epochs=epochs, validation_data=val_gen, callbacks=callbacks)
Explanation: Train the model
End of explanation
tf.saved_model.save(model, 'unet_pets_saved_model')
Explanation: Save the trained model in savedmodel format
End of explanation
def plot_results(x, y, preds):
    # x and y has a batch of images, we will visualize the first one
    idx = 1
    fig = plt.figure(figsize=(16,5))
    ax = fig.add_subplot(131)
    ax.imshow(x[idx,:,:,:].astype(np.uint8))
    ax.set_title('input')
    ax = fig.add_subplot(132)
    ax.imshow(y[idx,:,:,:])
    ax.set_title('ground_truth')
    ax = fig.add_subplot(133)
    # process the prediction
    mask = np.argmax(preds[idx], axis=-1)
    mask = np.expand_dims(mask, axis=-1)
    img = PIL.ImageOps.autocontrast(keras.preprocessing.image.array_to_img(mask))
    #ax.imshow(preds[idx,:,:,:])
    ax.imshow(img)
    ax.set_title('predictions')
    plt.show()
x, y = val_gen[0]
val_preds = model.predict(tf.constant(x))
plot_results(x, y, val_preds)
Explanation: Visualize predictions
End of explanation
def get_func_from_saved_model(saved_model_dir):
    saved_model_loaded = tf.saved_model.load(
        saved_model_dir, tags=[tag_constants.SERVING])
    graph_func = saved_model_loaded.signatures[
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    return graph_func, saved_model_loaded
def predict_and_benchmark_throughput(batched_input, model, N_warmup_run=50, N_run=500,
                                     result_key='predictions'):
    elapsed_time = []
    all_preds = []
    batch_size = batched_input.shape[0]
    elapsed_time = np.zeros(N_run)
    for i in range(N_warmup_run):                                             
        preds = model(batched_input)
    
    tmp = 0
    for i in range(N_run):
        start_time = time.time()
        preds = model(batched_input)
        # Force device synchronization with .numpy()
        tmp += preds[result_key][0,0].numpy() 
        end_time = time.time()
        elapsed_time[i] = end_time - start_time
        all_preds.append(preds)
        if i >= 50 and i % 50 == 0:
            print('Steps {}-{} average: {:4.1f}ms'.format(i-50, i, (elapsed_time[i-50:i].mean()) * 1000))
            
    print('Latency {:4.1f}+/-{:4.1f} ms'.format(elapsed_time.mean(), elapsed_time.std()))
    print('Throughput: {:.0f} images/s'.format(N_run * batch_size / elapsed_time.sum()))
    return all_preds
Explanation: Helper functions
End of explanation
func, _ = get_func_from_saved_model('unet_pets_saved_model')
output = func(tf.constant(x))
result_key = list(output.keys())[0]
output = output[result_key]
print('output shape', output.shape)
res = predict_and_benchmark_throughput(tf.constant(x), func, result_key=result_key)
Explanation: Run inference
End of explanation
img_size = (200, 200)
batch_size = 4
val_200 = OxfordPets(batch_size, img_size, val_inp_paths, val_target_paths)
x, y = val_200[0]
preds_200 = func(tf.constant(x))[result_key]
plot_results(x, y, preds_200)
Explanation: Now prepare input data with different resolution and run inference using that
End of explanation
def trt_convert(input_path, output_path, input_shapes, dynamic_shape=False, prof_strategy='Range'):
    conv_params=trt.TrtConversionParams(
        precision_mode='FP16', minimum_segment_size=30,
        max_workspace_size_bytes=1<<30, maximum_cached_engines=1)
    converter = trt.TrtGraphConverterV2(
        input_saved_model_dir=input_path, conversion_params=conv_params,
        use_dynamic_shape=dynamic_shape, 
        dynamic_shape_profile_strategy=prof_strategy)
    
    converter.convert()
    def input_fn():
        for shapes in input_shapes:
            # return a list of input tensors
            yield [np.ones(shape=x).astype(np.float32) for x in shapes]
    converter.build(input_fn)
    converter.save(output_path)
input_shapes = [[(4,160,160,3)], [(4,200,200,3)]]
trt_convert(input_path="unet_pets_saved_model", output_path="unet_pets_trt",
            input_shapes=input_shapes, dynamic_shape=True, prof_strategy='Range')
Explanation: We have seen so far that a fully convolutional TF model can be used to infer images with different input sizes.
Convert the model with TF-TRT
In dynamic shape mode we will have a single engine that can handle the input shapes seen during build mode. We select profile strategy 'Range', and we need to define the range of input shapes.
End of explanation
trt_func, _ = get_func_from_saved_model('unet_pets_trt')
x, y = val_gen[0]
trt_output = trt_func(tf.constant(x))[result_key]
plot_results(x, y, trt_output)
x, y = val_gen[10]
res = predict_and_benchmark_throughput(tf.constant(x), trt_func, result_key=result_key)
Explanation: Load converted model and run inference
End of explanation
x, y = val_200[10]
res = predict_and_benchmark_throughput(tf.constant(x), trt_func, result_key=result_key)
Explanation: Predict with different image size.
End of explanation
<END_TASK>