Running TRTorch on ResNet50 Torchvision model

TRTorch_resnet50.py

import trtorch
import torch
import torchvision.models as models


rn50 = models.resnet50(pretrained=True).cuda().eval()
inp = torch.randn((1, 3, 224, 224),  device='cuda')
rn50_jit_traced = torch.jit.trace(rn50, inp)
rn50_jit_scripted = torch.jit.script(rn50)

compile_settings = {
        "input_shapes" : [ (1, 3, 224, 224)],
        #"op_precision" : torch.half,
        "workspace_size" : 1<<30,
        }

trtorch_rn50_traced_model = trtorch.compile(rn50_jit_traced, compile_settings)
trtorch_rn50_scripted_model = trtorch.compile(rn50_jit_scripted, compile_settings)



rn50_output = rn50(inp)
rn50_trt_traced_output = trtorch_rn50_traced_model(inp)
rn50_trt_scripted_output = trtorch_rn50_scripted_model(inp)

print(torch.max(torch.abs(rn50_output - rn50_trt_traced_output)))
print(torch.max(torch.abs(rn50_output - rn50_trt_scripted_output)))
print(torch.max(torch.abs(rn50_trt_traced_output - rn50_trt_scripted_output)))

Steps to run:

1. Run NGC PyTorch Docker
   At the time of the script I ran docker run -it --gpus=all --ipc=host -v=<path>:/host nvcr.io/nvidia/pytorch:20.03-py3
2. Download a prebuilt TRTorch package
   wget https://github.com/NVIDIA/TRTorch/releases/download/v0.0.3/trtorch-0.0.3-cp36-cp36m-linux_x86_64.whl
3. Install the prebuilt package:
   pip install trtorch-0.0.3-cp36-cp36m-linux_x86_64.whl
4. Run the above script
   python TRTorch_resnet50.py