PyTorch 预训练模型转 onnx，Tensorrt 部署

所用环境

• Centos 7
• PyTorch 1.8.0（此为模型训练时的配置，转 onnx 时改为1.2.0）
• onnx 1.7.0
• tensorrt 7.2.2.3
• cuda 11.1

安装 onnx

pip install onnx

安装 tensorrt

另一位博主 @幸福回头 的博文中 (博文地址) 对于安装 tensorrt 有详细可靠的步骤，在此不作赘述。

PyTorch 预训练模型 .pth 文件到 onnx 文件的转换

这一步骤使用python完成，话不多说，上代码

import sys
import os 
sys.path.append(os.path.abspath(os.path.join(os.getcwd(), ".")))
import onnx 
import torch
from resnet50Pretrain import model_bn

model = model_bn(1024, 1226).cuda()

state_dict = torch.load("/search/gpu4/train_data/torch2onnx/res50_0.77.pth")['state_dict']
model.load_state_dict(state_dict) #预训练模型加载
model.cuda()

export_onnx_file = "resnet50_0728_77.onnx"
x=torch.onnx.export(model,  # 待转换的网络模型和参数
                torch.randn(1, 3, 300, 300, device='cuda'), # 虚拟的输入，用于确定输入尺寸和推理计算图每个节点的尺寸
                export_onnx_file,  # 输出文件的名称
                verbose=False,      # 是否以字符串的形式显示计算图
                input_names=["input"], #输入节点的名称，
                output_names=["output"], # 输出节点的名称
                opset_version=10,   # onnx 支持采用的operator set
                do_constant_folding=True, # 是否压缩常量
                dynamic_axes={"input":{0: "batch_size"}, "output":{0: "batch_size"},} #设置动态维度，此处指明input节点的第0维度可变，命名为batch_size，若设置固定的batch_size，则不需要设置dynamic_axes参数。
                )

net = onnx.load("resnet50_0728_77.onnx")  # 加载onnx 计算图
onnx.checker.check_model(net)  # 检查文件模型是否正确

import onnxruntime
import numpy as np

session = onnxruntime.InferenceSession("./resnet50_0728_77.onnx") # 创建一个运行session，类似于tensorflow
out_r = session.run(None, {"input": np.random.rand(1, 3, 300, 300).astype('float32')})  # 模型运行，注意这里的输入必须是numpy类型
print(len(out_r))
print(out_r[0].shape)

onnx 文件到 .trt 文件的转换

实现onnx 文件到 .trt 文件的转换不需要自己写转换文件，tensorrt 官方的安装包内已经提供了这样的文件，直接运行即可。找到 tensorrt 安装包位置，进入 trtexec 文件夹并 make：

cd TensorRT-7.2.2.3/samples/trtexec
make

此时在 TensorRT-7.2.2.3 文件夹内的 bin 文件夹下会生成一个可执行文件 trtexec， 可用于 onnx 到 trt 文件转换：

cd TensorRT-7.2.2.3/bin

若输入为固定 batch_size:

./trtexec --explicitBatch=1 --onnx=model.onnx --saveEngine=test/resnet50_0728_77.trt --int8 --workspace=1 --verbose

其中 “–int8” 选项为保存 .trt 文件时选择的精度为8位整形，若想要保存最高精度 fp32 ，去掉这个选项就可以。

若输入为动态 batch_size:

./trtexec  --onnx=torch2onnx/resnet50_0728_77.onnx --saveEngine=tt/resnet50_0728_77.trt --minShapes=input:1x3x300x300 --optShapes=input:16x3x300x300 --maxShapes=input:32x3x300x300 --shapes=input:1x3x300x300 --int8 --workspace=1 --verbose

到这里没有成功，产生 error ：Assertion failed: dims.nbDims == 4 || dims.nbDims == 5

[8] Assertion failed: dims.nbDims == 4 || dims.nbDims == 5
[07/28/2021-12:10:20] [E] Failed to parse onnx file
[07/28/2021-12:10:20] [E] Parsing model failed
[07/28/2021-12:10:20] [E] Engine creation failed
[07/28/2021-12:10:20] [E] Engine set up failed

经查阅资料，这里为版本错误，由于前面的 .pth 模型转换 onnx 文件时使用的 PyTorch 版本为1.8.0（与训练时 PyTorch 版本无关），需要降低到<=1.3.0 方可正常运行。重装版本为1.2.0后，运行成功：