Ubuntu下TensorRT的安装和torch2trt的配置及使用

目录

1、环境需求

2、安装tensorrt 7.2.2.3

​3、安装tensorrt7.2.2.3

4、下载torch2trt

5、设置torch2trt

6、运行torch2trt的demo

错误说明

torch2trt是NVIDIA官方给出的将pytorch模型转换为tensorrt引擎的工具。

1、环境需求

tensorrt 7.2.2.3（注意如果你的版本是tensorrt 8的，之后使用torch2trt会报错，具体之后我会说明）。

python = 3.8

至于cuda和cuDNN ，下载tensorrt 7.2.2.3的时候会自动下载相应的版本。

2、安装tensorrt 7.2.2.3

创建名为tensorrt的虚拟环境

conda create -n tensorrt python=3.8

安装pytorch框架

在pytorch官网找到相应的版本进行下载，指令如下：

pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

3、安装tensorrt7.2.2.3

参考了TensorRT 笔记 - 在 Conda 虚拟环境中安装 

pip install nvidia-pyindex
pip install nvidia-tensorrt==7.2.2.3

4、下载torch2trt

GitHub - NVIDIA-AI-IOT/torch2trt: An easy to use PyTorch to TensorRT converter

5、设置torch2trt

官方提供了3种设置，如图所示：

option1指的是没有插件，option2指的是有插件，插件的作用是支持使用pytorch中有的而tensorrt中没有的一些操作,option3支持实验性社区贡献的功能。这里我是使用的第1种设置。

注意：这里我试过用option2的方式设置，但是会出现如下图所示的错误：

官方的解释option2需要pytorch1.3+版本进行序列化，而我的pytorch版本是1.10的，因此就是用的option1.。

配置好之后，按照下图的方式测试torch2trt是否可用，如果不报错，便是torch2trt测试成功。

 

6、运行torch2trt的demo

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()

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

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

# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))
print(f"pytorch:{model}")
print(f"trt:{model_trt}")

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

结果：

错误说明

由于我之前使用的是tensorrt 8版本的，所以运行demo会报错

AttributeError: 'tensorrt.tensorrt.Builder' object has no attribute 'max_workspace_size'

官方的建议是使用tensorrt7.2版本的。AttributeError: 'tensorrt.tensorrt.Builder' object has no attribute 'max_workspace_size' · Issue #557 · NVIDIA-AI-IOT/torch2trt · GitHub

使用指令下载tensotensorrt7.2*版本的出错，参见第三步，安装tensorrt 7.2.3