trt | torch2trt的使用方式

关注

原创最新推荐文章于 2026-10-23 14:26:51 发布 1.8k 阅读 AI 写同款· GEO 优化 ›

快速部署使用tensorRT加速推理(trt,onnx) TensorRT nvidia推理框架，原理自行百度跑推理的时候用的两种方式第一种nvidia自家的转换这里拿pytorch做举例，其他的框架差不多 torch2trt 这里太简单了我就把官方例子给一下 import torch from torch2trt import torch2trt from torchvision.models.alexnet import alexnet # create some regular pytorch model... model = alexnet(p 阅读详情

一、安装

1. 安装 tensorrt python 接口

下载 trt 包 .tar.gz

https://developer.nvidia.com/nvidia-tensorrt-5x-download

解压

tar xvf TensorRT-6.0.1.5.Ubuntu-18.04.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz

安装 trt python 接口

cd python

pip install tensorrt-6.0.1.5-cp37-none-linux_x86_64.whl

安装 uff

cd uff

pip install uff-0.6.5-py2.py3-none-any.whl

验证trt是否安装成功

python

import tensorrt

2. 安装 torch2trt

sudo apt-get install libprotobuf* protobuf-compiler ninja-build

git clone https://github.com/NVIDIA-AI-IOT/torch2trt

cd torch2trt

sudo python setup.py install --plugins

二、代码演示

model = BNNproAtt()
model.load_state_dict(torch.load('/src/2_toeng/pytorch_2_eng/reid2trt/BNNproAtt0525.pt', map_location = 'cpu')
model.eval().float().cuda()

input_data = torch.rand((4, 3, 384, 128), dtype = torch.float).cuda()

t0 = time.time()
out = model(input_data)
t1 = time.time()
print("pytorch costed time: ", t1 - t0)

# convert to TensorRT model
model_trt = torch2trt(model, [input_data], max_batch_size = 4, int8_mode = True)

t2 = time.time()
out_trt = model_trt(input_data)
t3 = time.time()
print("trt costed time: ", t3 - t2)

# check the output against pytorch
print(torch.max(torch.abs(out - out_trt)

torch.save(model_trt.state_dict(), '/src/2_toeng/pytroch_2_eng/reid2trt/bnn_trt_int8.pt'

收录于