一、安装

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'