<jupyter_start><jupyter_text>TensorRT inference with torch2trt![process diagram](./doc/torchvision-torch2trt.svg)  IMPORTANT!**Before you start this notebook, please shutdown other notebook kernels to allocate enough memory for this notebook.**    Prepare model and input image Define a function to download a file.<jupyter_code>import urllib
import os

def download_file(url, file=None):
    if file is None:
        file = os.path.basename(url)
    if os.path.exists(file):
        return file
    try:
        urllib.URLopener().retrieve(url, file)
    except:
        urllib.request.urlretrieve(url, file)
    return file<jupyter_output><empty_output><jupyter_text>Specify a input image file.<jupyter_code>image_file = 'bellpepper.jpg'
#image_file = download_file('https://github.com/pytorch/hub/raw/master/dog.jpg')
#image_file = download_file('https://upload.wikimedia.org/wikipedia/commons/b/b5/Canadian_Golden_Retriever.jpeg')<jupyter_output><empty_output><jupyter_text>Download the ImageNet class file.<jupyter_code>class_file = download_file('https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json')<jupyter_output><empty_output><jupyter_text>Parse the class file and convert it to a Python dictionary.<jupyter_code>import json
class_dict = json.load(open(class_file, 'r'))<jupyter_output><empty_output><jupyter_text>Export the input image file to PyTorch tensor and normlize it in the ImageNet training image conditions.<jupyter_code># sample execution (requires torchvision)
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np

input_image = Image.open(image_file)
plt.imshow(input_image)
plt.show()

preprocess = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model

filtered_image = input_tensor.clone()
filtered_image = filtered_image.numpy().transpose((1, 2, 0))
filtered_image = np.float32(filtered_image)
filtered_image = np.clip(filtered_image, 0, 1)
plt.imshow(filtered_image)
plt.show()<jupyter_output><empty_output><jupyter_text>Download the Torchvision AlexNet pre-traind model.<jupyter_code>from torchvision import models
net = models.alexnet(pretrained=True)
print(net)<jupyter_output>AlexNet(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))
    (1): ReLU(inplace=True)
    (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (4): ReLU(inplace=True)
    (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (7): ReLU(inplace=True)
    (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (9): ReLU(inplace=True)
    (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(6, 6))
  (classifier): Sequential(
    (0): Dropout(p=0.5, inplace=False)
    (1): Linear(in_features=9216, out_features=4096, bias=True)
 [...]<jupyter_text>Export the PyTorch model to TensorRT with torch2trt<jupyter_code>import torch
from torch2trt import torch2trt

x = torch.zeros(1, 3, 224, 224).cuda()
model_trt = torch2trt(net.cuda(), [x])<jupyter_output><empty_output><jupyter_text>Do inference by TensorRT with torch2trt<jupyter_code>y_trt = model_trt(input_batch.cuda()).cpu()
outputs = y_trt.numpy()
idxlist = np.argsort(outputs[0])[:-4:-1]
for i in idxlist:
    print('{}:\t{}'.format(class_dict[str(i)][1], outputs[0][i]))<jupyter_output>bell_pepper:	15.869132995605469
cucumber:	13.745673179626465
Granny_Smith:	12.043103218078613<jupyter_text>Measure inference processing time Define a decorator function to measure processing times.<jupyter_code>import time
import functools

def processing_time(num_times, warmups=0):
    def _processing_time(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for i in range(warmups):
                ret = func(*args, **kwargs)
            t0 = time.time()
            for i in range(num_times):
                ret = func(*args, **kwargs)
            t1 = time.time()
            print('Average a {0} times trial: {1:,f}sec'.format(num_times, (t1 - t0) / num_times))
            return ret
        return wrapper
    return _processing_time<jupyter_output><empty_output><jupyter_text>Define a test function with the pocessing time measurement decorator.<jupyter_code>@processing_time(100, 10)
def test_inference(network, in_dat):
    out_dat = network(in_dat)
    return out_dat<jupyter_output><empty_output><jupyter_text>Measure inference processing time.<jupyter_code>image_cuda = input_batch.cuda()
torch.cuda.current_stream().synchronize()
y_trt = test_inference(model_trt, image_cuda).cpu()
torch.cuda.current_stream().synchronize()
outputs = y_trt.numpy()
idxlist = np.argsort(outputs[0])[:-4:-1]
for i in idxlist:
    print('{}:\t{}'.format(class_dict[str(i)][1], outputs[0][i]))<jupyter_output>Average a 100 times trial: 0.015197sec
bell_pepper:	15.869132995605469
cucumber:	13.745673179626465
Granny_Smith:	12.043103218078613<jupyter_text>Print the TensorRT network.<jupyter_code>import trt_analyzer
import pandas as pd

net_dict = trt_analyzer.network_dict(model_trt.network)
pd.set_option('max_colwidth', 1000)
pd.DataFrame(net_dict)<jupyter_output><empty_output><jupyter_text>Display a diagram of the TensorRT network.<jupyter_code>from torch2trt.utils import trt_network_to_dot_graph

dot = trt_network_to_dot_graph(model_trt.network)
dot<jupyter_output><empty_output><jupyter_text>Save the TensorRT model.  **Jetson Nano may fail this process for lack of memory.**<jupyter_code>'''
model_file_path = 'model_torch2trt.pth'
torch.save(model_trt.state_dict(), model_file_path)
'''<jupyter_output><empty_output><jupyter_text>Load the saved TensorRT model and do inference with it.  **Jetson Nano may fail this process for lack of memory.**<jupyter_code>'''
from torch2trt import TRTModule

model_trt = TRTModule()
model_trt.load_state_dict(torch.load(model_file_path))

y_trt = model_trt(input_batch.cuda()).cpu()
outputs = y_trt.numpy()
idxlist = np.argsort(outputs[0])[:-4:-1]
for i in idxlist:
    print('{}:\t{}'.format(class_dict[str(i)][1], outputs[0][i]))
'''<jupyter_output><empty_output><jupyter_text>Work with custom layer Import the pooling custom layer handling modules<jupyter_code>import sys

cur_path = %pwd
plugin_path = os.path.join(cur_path, 'plugin')
sys.path.append(plugin_path)
from trt_plugin_pb2 import copy_Message
from trt_plugin_pb2 import pooling_Message
import trt_plugin_pb2<jupyter_output><empty_output><jupyter_text>Load the pooling custom layer plugin.<jupyter_code>import ctypes

lib_file = os.path.join(plugin_path, 'build', 'libPoolingPlugin.so')
lib = ctypes.CDLL(lib_file)<jupyter_output><empty_output><jupyter_text>Get the plugin registry.<jupyter_code>import tensorrt as trt

registry = trt.get_plugin_registry()
print([c.name for c in registry.plugin_creator_list])
print([c.plugin_namespace for c in registry.plugin_creator_list])<jupyter_output>['RnRes2Br2bBr2c_TRT', 'RnRes2Br2bBr2c_TRT', 'RnRes2Br1Br2c_TRT', 'RnRes2Br1Br2c_TRT', 'CustomSkipLayerNormPluginDynamic', 'CustomEmbLayerNormPluginDynamic', 'CustomGeluPluginDynamic', 'CustomQKVToContextPluginDynamic', 'CustomFCPluginDynamic', 'SingleStepLSTMPlugin', 'interpolate', 'interpolate', 'pooling', 'copy']
['', '', '', '', '', '', '', '', '', '', 'torch2trt', 'torch2trt', 'macnica_trt_plugins', 'macnica_trt_plugins']<jupyter_text>Register the plugin creator.<jupyter_code>namespace = 'macnica_trt_plugins'
macnica_creators = [c for c in registry.plugin_creator_list if c.plugin_namespace == namespace]
for c in macnica_creators:
    registry.register_creator(c, namespace)<jupyter_output><empty_output><jupyter_text>Define a plugin converter.<jupyter_code>from torch2trt.torch2trt import *

@tensorrt_converter('torch.nn.functional.max_pool2d')
def convert_max_pool2d_plugin(ctx):
    # parse args
    input = get_arg(ctx, 'input', pos=0, default=None)
    kernel_size = get_arg(ctx, 'kernel_size', pos=1, default=None)
    stride = get_arg(ctx, 'stride', pos=2, default=None)
    padding = get_arg(ctx, 'padding', pos=3, default=0)
    dilation = get_arg(ctx, 'dilation', pos=4, default=1)
    ceil_mode = get_arg(ctx, 'ceil_mode', pos=5, default=False)
    
    # get input trt tensor (or create constant if it doesn't exist)
    input_trt = trt_(ctx.network, input)
    
    output = ctx.method_return

    # get kernel size
    if not isinstance(kernel_size, tuple):
        kernel_size = (kernel_size, ) * 2

    # get stride
    if not isinstance(stride, tuple):
        stride = (stride, ) * 2

    # get padding
    if not isinstance(padding, tuple):
        padding = (padding, ) * 2
    
    print('convert_max_pool2d_plugin:')
    print('input={}, kernel_size={}, stride={}, padding={}'.format( \
        input_trt.shape, kernel_size, stride, padding))
    
    registry = trt.get_plugin_registry()
    creator = registry.get_plugin_creator( \
        type='pooling', version='1', plugin_namespace='macnica_trt_plugins')
    
    sz = input_trt.shape
    if len(sz) == 3:
        sz = (1, sz[0], sz[1], sz[2])
    
    impl = trt_plugin_pb2.CUDA
    #impl = trt_plugin_pb2.CuDNN
    message = pooling_Message( \
        dims=sz, mode=trt_plugin_pb2.Maximum, window=kernel_size, stride=stride, impl=impl)
    plg = creator.deserialize_plugin('pooling', message.SerializeToString())
    layer = ctx.network.add_plugin_v2(inputs=[input_trt], plugin=plg)

    output._trt = layer.get_output(0)<jupyter_output><empty_output><jupyter_text>Generate a TensorRT network with the plugin.<jupyter_code>x = torch.zeros(1, 3, 224, 224).cuda()
model_trt = torch2trt(net.cuda(), [x])<jupyter_output>convert_max_pool2d_plugin:
input=(64, 55, 55), kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)
convert_max_pool2d_plugin:
input=(192, 27, 27), kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)
convert_max_pool2d_plugin:
input=(256, 13, 13), kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)<jupyter_text>Print the TensorRT network.<jupyter_code>net_dict = trt_analyzer.network_dict(model_trt.network)
pd.DataFrame(net_dict)<jupyter_output><empty_output><jupyter_text>Do inference by TensorRT with the generated network.<jupyter_code>y_trt = model_trt(input_batch.cuda()).cpu()
outputs = y_trt.numpy()
idxlist = np.argsort(outputs[0])[:-4:-1]
for i in idxlist:
    print('{}:\t{}'.format(class_dict[str(i)][1], outputs[0][i]))<jupyter_output>bell_pepper:	15.869132995605469
cucumber:	13.745673179626465
Granny_Smith:	12.043103218078613
{'fraction_non_alphanumeric': 0.13420030705319244, 'fraction_numerical': 0.041542490743249345, 'mean_word_length': 5.8442521631644, 'pattern_counts': {'":': 0, '<': 87, '<?xml version=': 0, '>': 87, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '23', 'dir': 'jupyter-structured-clean-dedup', 'id': '447616', 'license': 'MIT', 'n_tokens_mistral': 4302, 'n_tokens_neox': 4028, 'n_words': 782, 'path': 'torchvision_torch2trt.ipynb', 'repo_name': 'MACNICA-CLAVIS-NV/torchvision2trt-samples', 'size': '11073'}