git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install
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'))
Original: https://blog.csdn.net/dou3516/article/details/124538557
Author: there2belief
Title: 【TensorRT】torch2trt安装与使用简例
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/688708/
转载文章受原作者版权保护。转载请注明原作者出处！