Basic Usage

This page demonstrates basic torch2trt usage.

Conversion

You can easily convert a PyTorch module by calling torch2trt passing example data as input, for example to convert alexnet we call

import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()

# create example data
x = torch.ones((1, 3, 224, 224)).cuda()

# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])

Note

Currently with torch2trt, once the model is converted, you must use the same input shapes during execution. The exception is the batch size, which can vary up to the value specified by the max_batch_size parameter.

Executution

We can execute the returned TRTModule just like the original PyTorch model. Here we execute the model and print the maximum absolute error.

y = model(x)
y_trt = model_trt(x)

# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))

Saving and loading

We can save the model as a state_dict.

torch.save(model_trt.state_dict(), 'alexnet_trt.pth')

We can load the saved model into a TRTModule

from torch2trt import TRTModule

model_trt = TRTModule()

model_trt.load_state_dict(torch.load('alexnet_trt.pth'))