序言
这篇文章总结了训练好的pytorch模型转成tensorrt模型部署的几种方式，转换原理流程大致如下：
- 导出网络定义以及相关权重；
- 解析网络定义以及相关权重；
- 根据显卡算子构造出最优执行计划；
- 将执行计划序列化存储；
- 反序列化执行计划；
- 进行推理
值得注意的是第三点，可以看到tensorrt转换出来的模型实际上是和硬件绑定的，也就是在部署的过程中，如果你的显卡和显卡相关驱动软件（cuda、cudnn）发生了改变，那么模型就得需要重新做转换。
一、trtexec
trtexec是在tensorrt包中自带的转换程序，该程序位于bin目录下，用起来比较方便，也是最简单的trt模型转换方式，在使用之前需要系统安装好cuda和cudnn，否则无法正常运行。使用示例如下：
首先将pytorch模型先转换成onnx模型，示例代码如下：
def torch2onnx(model_path,onnx_path):
model = load_model(model_path)
test_arr = torch.randn(1,3,32,448)
input_names = ['input']
output_names = ['output']
tr_onnx.export(
model,
test_arr,
onnx_path,
verbose=False,
opset_version=11,
input_names=input_names,
output_names=output_names,
dynamic_axes={"input":{3:"width"}} #动态推理W纬度，若需其他动态纬度可以自行修改，不需要动态推理的话可以注释这行
)
print('->>模型转换成功！')
trtexec转换命令如下：
固定尺寸模型转换：
./trtexec --onnx=repvgg_a1.onnx --saveEngine=repvgg_a1.engine --workspace=1024 --fp16
动态尺寸模型转换：
./trtexec --onnx=repvgg_a1.onnx --saveEngine=repvgg_a1.engine --workspace=1024 --minShapes=input:1x3x32x32 --optShapes=input:1x3x32x320 --maxShapes=input:1x3x32x640 --fp16
参数详解：
- –onnx onnx路径
- –saveEngine trt序列化推理引擎保存地址
- –workspace 以兆字节为单位设置工作区大小(默认= 16)
- –minShapes 使用提供的最小形状的配置文件生成动态形状
- –optShapes 使用提供的最优形状的配置文件生成动态形状
- –maxShapes 使用提供的最大形状的配置文件生成动态形状
- –fp16 开启 float16精度的推理（推荐此模式，一方面能够加速，另一方面精度下降比较小）
二、torch2trt
torch2trt是nvidia官方维护的一个易于使用的PyTorch到TensorRT转换器，使用起来也比较简单，但是相对于上面的方式环境配置更复杂一些，需要预先安装好torch、torch2trt、tensorrt，python环境下tensorrt安装方式示例，在Tensorrt的 .tar包中依次找到这几个whl包，直接使用pip安装即可：
#1、安装tensorrt
cd ~/TensorRT-8.2.4.2/python
pip install tensorrt-8.2.4.2-cp37-none-linux_x86_64.whl
#2、安装Python UFF wheel文件。只有当你将TensorRT与TensorFlow一起使用时才需要安装这个文件 用处：pb转tensorRT
cd ~/TensorRT-8.2.4.2/uff
pip install uff-0.6.9-py2.py3-none-any.whl
#3、安装Python graphsurgeon whl文件 用处：可以让TensorRT 自定义网络结构
cd ~/TensorRT-8.2.4.2/graphsurgeon
pip install graphsurgeon-0.4.5-py2.py3-none-any.whl
#注意trt7.0的版本没有这个包(不用装)
#4、安装Python onnx-graphsurgeon whl文件
cd ~/TensorRT-8.2.4.2/onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl
#5、安装pycuda 可以通过它来实现python 下CUDA 的编程
pip install pycuda
#6、验证安装,打印出tensorrt版本，即安装成功
python
import tensorrt
tensorrt.__version__
torch2trt安装：
git clone https://github.com/NVIDIA-AI-IOT/torch2trt.git
cd torch2trt
sudo python setup.py install --plugins
模型转换代码使用示例：
model = load_model(model_path)
model.cuda()
arr = torch.ones(1, 3, 32, 448).cuda()
model_trt = torch2trt(model,
[arr],
fp16_mode=True,
log_level=trt.Logger.INFO,
max_workspace_size=(1 << 32),
max_batch_size=1,
)
torch.save(model_trt.state_dict(), os.path.join(output, "model_trt.pth"))
logger.info("Converted TensorRT model done.")
engine_file = os.path.join(output, "model_trt.engine")
with open(engine_file, "wb") as f:
f.write(model_trt.engine.serialize())
logger.info("Converted TensorRT model engine file is saved for C++ inference.")
其中保存的pth模型可以用来加载到TRTModule中,TRTModule可以像torch模型一样正常推理：
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('model_trt.pth'))
而engine序列化文件可用于tensorrt程序加载推理，但是这里需要注意的是torch2trt不支持动态推理，更多使用示例参考torch2trt的github说明。
三、torch2trt dynamic
torch2trt dynamic是torch2trt的动态推理版，支持模型转换后的动态推理，使用起来和torch2trt基本一致，首先是安装：
git clone https://github.com/grimoire/torch2trt_dynamic.git
cd torch2trt_dynamic
python setup.py develop
然后是使用示例，增加了一个动态尺度的参数：
from torch2trt_dynamic import torch2trt_dynamic
import torch
from torch import nn
from torchvision.models.resnet import resnet50
import os
# create some regular pytorch model...
model = resnet50().cuda().eval()
# create example data
x = torch.ones((1, 3, 224, 224)).cuda()
# convert to TensorRT feeding sample data as input
opt_shape_param = [
[
[1, 3, 128, 128], # min
[1, 3, 256, 256], # opt
[1, 3, 512, 512] # max
]
]
model_trt = torch2trt_dynamic(model, [x], fp16_mode=False, opt_shape_param=opt_shape_param)
torch.save(model_trt.state_dict(), os.path.join(output, "model_trt.pth"))
logger.info("Converted TensorRT model done.")
engine_file = os.path.join(output, "model_trt.engine")
with open(engine_file, "wb") as f:
f.write(model_trt.engine.serialize())
logger.info("Converted TensorRT model engine file is saved for C++ inference.")
四、手动解析onnx模型
如果不想借助工具转换，也可以自己编写代码，使用tensorrt的接口手动解析onnx模型，构建engine引擎，这种方式比较简单，不依赖其他库，并且支持动态推理模型转换，代码示例如下：
# --*-- coding:utf-8 --*--
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import time
import cv2, os
import numpy as np
import math
TRT_LOGGER = trt.Logger()
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
"""
host_mem: cpu memory
device_mem: gpu memory
"""
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="", fp16_mode=False, save_engine=False,input_dynamic=False):
"""
params max_batch_size: 预先指定大小好分配显存
params onnx_file_path: onnx文件路径
params engine_file_path: 待保存的序列化的引擎文件路径
params fp16_mode: 是否采用FP16
params save_engine: 是否保存引擎
returns: ICudaEngine
"""
# 如果已经存在序列化之后的引擎，则直接反序列化得到cudaEngine
if os.path.exists(engine_file_path):
print("Reading engine from file: {}".format(engine_file_path))
with open(engine_file_path, 'rb') as f, \
trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read()) # 反序列化
else: # 由onnx创建cudaEngine
# 使用logger创建一个builder
# builder创建一个计算图 INetworkDefinition
explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
# In TensorRT 7.0, the ONNX parser only supports full-dimensions mode, meaning that your network definition must be created with the explicitBatch flag set. For more information, see Working With Dynamic Shapes.
with trt.Builder(TRT_LOGGER) as builder, \
builder.create_network(explicit_batch) as network, \
trt.OnnxParser(network, TRT_LOGGER) as parser: # 使用onnx的解析器绑定计算图，后续将通过解析填充计算图
# builder.max_workspace_size = 1 << 30 # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30
builder.max_batch_size = max_batch_size # 执行时最大可以使用的batchsize
if fp16_mode:
config.set_flag(trt.BuilderFlag.FP16)
# builder.fp16_mode = fp16_mode
# 解析onnx文件，填充计算图
if not os.path.exists(onnx_file_path):
quit("ONNX file {} not found!".format(onnx_file_path))
print('loading onnx file from path {} ...'.format(onnx_file_path))
with open(onnx_file_path, 'rb') as model: # 二值化的网络结果和参数
print("Begining onnx file parsing")
parser.parse(model.read()) # 解析onnx文件
# parser.parse_from_file(onnx_file_path) # parser还有一个从文件解析onnx的方法
print("Completed parsing of onnx file")
# 填充计算图完成后，则使用builder从计算图中创建CudaEngine
print("Building an engine from file{}' this may take a while...".format(onnx_file_path))
if input_dynamic: # 动态推理
profile = builder.create_optimization_profile()
profile.set_shape("input",(1,3,32,32),(1,3,32,320),(1,3,32,640))
config.add_optimization_profile(profile)
#################
print(network.get_layer(network.num_layers - 1).get_output(0).shape)
engine = builder.build_engine(network, config)
print("Completed creating Engine")
if save_engine: # 保存engine供以后直接反序列化使用
with open(engine_file_path, 'wb') as f:
f.write(engine.serialize()) # 序列化
return engine
if __name__== "__main__":
# These two modes are depend on hardwares
fp16_mode = True
max_batch_size = 1
onnx_model_path = "./repvgg_a1.onnx"
trt_engine_path = "./repvgg_a1.engine"
# Build an cudaEngine
engine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode,True,True)
五、tensorrtx
tensorrtx的模型构建方式比较奇特，首先使用tensorrt自带API将网络搭建起来，然后再将权重赋值进去，这种方式在转换的时候只要网络搭得没问题，基本上转换也不会问题，很好的解决了在onnx转trt过程中某些算子不支持的问题，但是过程比较复杂，况且不支持动态尺度推理，目前trt对onnx的支持很好，基本上onnx的模型都能转过去，所以这种方式如果不嫌麻烦的话可以试试，示例过程：
克隆tensorrtx
git clone https://github.com/wang-xinyu/tensorrtx.git
生成yolov5.wts文件，下载权重文件yolov5s.pt，将tensorrtx/yolov5/gen_wts.py复制到ultralytics/yolov5中，执行
python gen_wts.py
编译tensorrtx/yolov5,生成yolov5.engine文件
mkdir build
cd build
cmake ..
make
默认情况下，生成的是s模型和fp16推理。以及批次为1的engine，yolov5的其他模型可以在代码中修改相关参数，
#define USE_FP16
#define DEVICE 0 // GPU ID
#define NMS_THRESH 0.4
#define CONF_THRESH 0.5
#define BATCH_SIZE 1
#define NET s // s m x l
copy文件yolov5.wts到tensorrtx/yolov5/build目录下，执行以下命令生成yolov5.engine
sudo ./yolov5 -s yolov5s.wts yolov5.engine s
sudo ./yolov5 -d yolov5s.engine ../samples
六、onnx-tensorrt
onnx-tensorrt是onnx官方的一个转换仓库，提供了诸多对应的Tensorrt版本的分支，比如我这里用的是8.2-EA，正确的编译方式为：
先将onnx-tensorrt git下来：
git clone --recursive -b 8.2-EA https://github.com/onnx/onnx-tensorrt.git
编译：
cd onnx-tensorrt
mkdir build
cd build
# /path/to/TensorRT-8.2.4.2改成自己的TensorRT绝对路径
cmake .. -DTENSORRT_ROOT=/path/to/TensorRT-8.2.4.2
make -j8
make install
编译完成后，如果你的cuda环境变量是配置好的，则不需要再配置，如果没有配置好则需要配置一下：
终端中输入：vim ~/.bashrc
#cuda
export PATH=/usr/local/cuda-11.4/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
保存退出，运行source ~/.bashrc刷新生效
onnx-tensorrt转换命令如下，序列化成engine引擎：
onnx2trt my_model.onnx -o my_engine.trt
转换为可读txt文本：
onnx2trt my_model.onnx -t my_model.onnx.txt
python端使用onnx-tensorrt：
#安装tensorrt
python3 -m pip install <tensorrt_install_dir>/python/tensorrt-8.x.x.x-cp<python_ver>-none-linux_x86_64.whl
#安装onnx
python3 -m pip install onnx==1.8.0
#安装onnx-tensorrt，在onnx-tensorrt目录下运行
python3 setup.py install
python代码使用示例，用来推理：
import onnx
import onnx_tensorrt.backend as backend
import numpy as np
model = onnx.load("/path/to/model.onnx")
engine = backend.prepare(model, device='CUDA:0')
input_data = np.random.random(size=(32, 3, 224, 224)).astype(np.float32)
output_data = engine.run(input_data)[0]
print(output_data)
print(output_data.shape)
