#!/usr/bin/env python
# coding: utf-8

# In[1]:


# Copyright 2022 NVIDIA Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================


# <img src="http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_dynamic_shapes/nvidia_logo.png" width="90px">
# 
# # Torch-TensorRT - Using Dynamic Shapes

# Torch-TensorRT is a compiler for PyTorch/TorchScript, targeting NVIDIA GPUs via NVIDIA's TensorRT Deep Learning Optimizer and Runtime. Unlike PyTorch's Just-In-Time (JIT) compiler, Torch-TensorRT is an Ahead-of-Time (AOT) compiler, meaning that before you deploy your TorchScript code, you go through an explicit compile step to convert a standard TorchScript program into a module targeting a TensorRT engine. Torch-TensorRT operates as a PyTorch extension and compiles modules that integrate into the JIT runtime seamlessly. After compilation, using the optimized graph should feel no different than running a TorchScript module. You also have access to TensorRT's suite of configurations at compile-time, so you are able to specify operating precision (FP32/FP16/INT8) and other settings for your module.
# 
# We highly encourage users to run this notebook using our NVIDIA's [PyTorch container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch). It comes packaged with a host of NVIDIA libraries and optimizations to widely used third-party libraries. In addition, this container is tested and updated on a monthly cadence!
# 
# This notebook has the following sections:
# 1. [TL;DR Explanation](#1)
# 2. [Setting up the model](#2)
# 3. [Working with Dynamic shapes in Torch TRT](#3)

# <a id="1"></a>
# ---
# ## TL;DR Explanation
# 
# Making use of Dynamic Shaped Tensors in Torch TensorRT is quite simple. Let's say you are using the `torch_tensorrt.compile(...)` function to compile a torchscript module.
# 
# One of the `args` in this function in this function is `input`: which defines an input to a module in terms of expected shape, data type and tensor format: `torch_tensorrt.Input`. 
# 
# For the purposes of this walkthrough we just need three `kwargs`: `min_shape`, `opt_shape` and `max_shape`. 
# ```
# ...
# torch_tensorrt.Input(
#         min_shape=(1, 224, 224, 3),
#         opt_shape=(1, 512, 512, 3),
#         max_shape=(1, 1024, 1024, 3),
#         dtype=torch.int32
#         format=torch.channel_last
#     )
# ...
# ```
# In this example, we are going to use a simple ResNet model to demonstrate the use of the API. We will be using different batch sizes in the example, but you can use the same method to alter any of the dimensions of the tensor.

# In[2]:


get_ipython().system('nvidia-smi')
get_ipython().system('pip install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org')


# <a id="2"></a>
# ---
# ## Setting up the model
# 
# In this section, we will:
# * Get sample data.
# * Download model from torch hub.
# * Build simple utility functions

# ### Getting sample data

# In[3]:


get_ipython().system('mkdir -p ./data')
get_ipython().system('wget  -O ./data/img0.JPG "https://d17fnq9dkz9hgj.cloudfront.net/breed-uploads/2018/08/siberian-husky-detail.jpg?bust=1535566590&width=630"')
get_ipython().system('wget  -O ./data/img1.JPG "https://www.hakaimagazine.com/wp-content/uploads/header-gulf-birds.jpg"')
get_ipython().system('wget  -O ./data/img2.JPG "https://www.artis.nl/media/filer_public_thumbnails/filer_public/00/f1/00f1b6db-fbed-4fef-9ab0-84e944ff11f8/chimpansee_amber_r_1920x1080.jpg__1920x1080_q85_subject_location-923%2C365_subsampling-2.jpg"')
get_ipython().system('wget  -O ./data/img3.JPG "https://www.familyhandyman.com/wp-content/uploads/2018/09/How-to-Avoid-Snakes-Slithering-Up-Your-Toilet-shutterstock_780480850.jpg"')

get_ipython().system('wget  -O ./data/imagenet_class_index.json "https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json"')


# In[4]:


# visualizing the downloaded images

from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
import json

fig, axes = plt.subplots(nrows=2, ncols=2)

for i in range(4):
    img_path = './data/img%d.JPG'%i
    img = Image.open(img_path)
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    input_tensor = preprocess(img)
    plt.subplot(2,2,i+1)
    plt.imshow(img)
    plt.axis('off')

# loading labels
with open("./data/imagenet_class_index.json") as json_file:
    d = json.load(json_file)


# ### Download model from torch hub.

# In[5]:


import torch

torch.hub._validate_not_a_forked_repo=lambda a,b,c: True

resnet50_model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
resnet50_model.eval()


# ### Build simple utility functions

# In[6]:


import numpy as np
import time
import torch.backends.cudnn as cudnn
cudnn.benchmark = True

def rn50_preprocess():
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return preprocess

# decode the results into ([predicted class, description], probability)
def predict(img_path, model):
    img = Image.open(img_path)
    preprocess = rn50_preprocess()
    input_tensor = preprocess(img)
    input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model

    # move the input and model to GPU for speed if available
    if torch.cuda.is_available():
        input_batch = input_batch.to('cuda')
        model.to('cuda')

    with torch.no_grad():
        output = model(input_batch)
        # Tensor of shape 1000, with confidence scores over Imagenet's 1000 classes
        sm_output = torch.nn.functional.softmax(output[0], dim=0)

    ind = torch.argmax(sm_output)
    return d[str(ind.item())], sm_output[ind] #([predicted class, description], probability)

# benchmarking models
def benchmark(model, input_shape=(1024, 1, 224, 224), dtype='fp32', nwarmup=50, nruns=10000):
    input_data = torch.randn(input_shape)
    input_data = input_data.to("cuda")
    if dtype=='fp16':
        input_data = input_data.half()

    print("Warm up ...")
    with torch.no_grad():
        for _ in range(nwarmup):
            features = model(input_data)
    torch.cuda.synchronize()
    print("Start timing ...")
    timings = []
    with torch.no_grad():
        for i in range(1, nruns+1):
            start_time = time.time()
            features = model(input_data)
            torch.cuda.synchronize()
            end_time = time.time()
            timings.append(end_time - start_time)
            if i%10==0:
                print('Iteration %d/%d, ave batch time %.2f ms'%(i, nruns, np.mean(timings)*1000))
                print('Images processed per second=', int(1000*input_shape[0]/(np.mean(timings)*1000)))
    print("Input shape:", input_data.size())
    print("Output features size:", features.size())
    print('Average batch time: %.2f ms'%(np.mean(timings)*1000))


# Let's test our util functions on the model we have set up, starting with simple predictions.

# In[7]:


for i in range(4):
    img_path = './data/img%d.JPG'%i
    img = Image.open(img_path)

    pred, prob = predict(img_path, resnet50_model)
    print('{} - Predicted: {}, Probablility: {}'.format(img_path, pred, prob))

    plt.subplot(2,2,i+1)
    plt.imshow(img)
    plt.axis('off')
    plt.title(pred[1])


# Onwards, to benchmarking.

# In[8]:


# Model benchmark without Torch-TensorRT
model = resnet50_model.eval().to("cuda")
benchmark(model, input_shape=(16, 3, 224, 224), nruns=100)


# ---
# ## Benchmarking with Torch-TRT (without dynamic shapes)

# In[9]:


import torch_tensorrt

trt_model_without_ds = torch_tensorrt.compile(model, inputs = [torch_tensorrt.Input((32, 3, 224, 224), dtype=torch.float32)],
    enabled_precisions = torch.float32, # Run with FP32
    workspace_size = 1 << 33
)


# In[10]:


benchmark(trt_model_without_ds, input_shape=(32, 3, 224, 224), nruns=100)


# With the baseline ready, we can proceed to the section working discussing dynamic shapes!

# <a id="3"></a>
# ---
# ## Working with Dynamic shapes in Torch TRT
#  
# Enabling "Dynamic Shaped" tensors to be used is essentially enabling the ability to defer defining the shape of tensors until run-time. Torch TensorRT simply leverages TensorRT's Dynamic shape support. You can read more about TensorRT's implementation in the [TensorRT Documentation](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#work_dynamic_shapes).
#  
# #### How can you use this feature?
#  
# To make use of dynamic shapes, you need to provide three shapes:
# * `min_shape`: The minimum size of the tensor considered for optimizations.
# * `opt_shape`: The optimizations will be done in an effort to maximize performance for this shape.
# * `max_shape`: The maximum size of the tensor considered for optimizations.
#  
# Generally, users can expect the best performance within the specified ranges. Performance for other shapes maybe be lower for other shapes (depending on the model ops and GPU used)
#  
# In the following example, we will showcase varying batch sizes, which is the zeroth dimension of our input tensors. As Convolution operations require that the channel dimension be a build-time constant, we won't be changing the sizes of other channels in this example, but for models which contain ops conducive to changes in other channels, this functionality can be freely used.

# In[11]:


# The compiled module will have precision as specified by "op_precision".
# Here, it will have FP32 precision.
trt_model_with_ds = torch_tensorrt.compile(model, inputs = [torch_tensorrt.Input(
        min_shape=(16, 3, 224, 224),
        opt_shape=(32, 3, 224, 224),
        max_shape=(64, 3, 224, 224),
        dtype=torch.float32)],
    enabled_precisions = torch.float32, # Run with FP32
    workspace_size = 1 << 33
)


# In[12]:


benchmark(trt_model_with_ds, input_shape=(16, 3, 224, 224), nruns=100)


# In[13]:


benchmark(trt_model_with_ds, input_shape=(32, 3, 224, 224), nruns=100)


# In[14]:


benchmark(trt_model_with_ds, input_shape=(64, 3, 224, 224), nruns=100)


# ## What's Next?
# 
# Check out the [TensorRT Getting started page](https://developer.nvidia.com/tensorrt-getting-started) for more tutorials, or visit the Torch-TensorRT [documentation](https://nvidia.github.io/Torch-TensorRT/) for more information!