pytorch模型转onnx再转tensorrt

前提条件：

pytorch模型文件（xxx.pt）

export_onnx.py

依赖库：

tensorrt 7.1.3

一、将torch模型转化为onnx模型

# 加载必要的依赖包
from modeling.deeplab import * 
import torch 
import torchvision 
import onnx 

# 定义加载原始模型的方法
def get_model(): 
    model = DeepLab(num_classes=4, backbone='mobilenet', output_stride=16,
    sync_bn=None, freeze_bn=False) 
    ckpt = torch.load('model_best.pth', map_location='cpu') 
    model.load_state_dict(ckpt['state_dict']) 
    model = model.cuda() 
    model.eval() 
    return model 

# 加载模型
model = get_model() 

# 配置onnx模型的节点名称，参数
input_names = ["input"] 
output_names = ["output"] 
example_tensor = torch.randn(1, 3, 320, 640, device='cuda') 
torch.onnx.export(model, example_tensor, "deeplab.onnx", verbose=True, 
opset_version=11, input_names=input_names, output_names=output_names) 

# 加载onnx模型 
onnx_model = onnx.load('deeplab.onnx') 

# 检查onnx模型是否正确导出
onnx.checker.check_model(onnx_model) 

二、使用onnx模型转化tensorrt模型

1.安装完tensorrt之后，使用tensorrt的工具trtexec，导出trt模型。

#该命令不会修改模型的精度类型 
sudo ./trtexec --onnx=alexnet.onnx --saveEngine=alexnet.trt 
#该命令会修改模型的精度类型为int8
sudo ./trtexec --onnx=alexnet.onnx --saveEngine=alexnet.trt --int8 

三、前后处理的重写

根据原算法的图像预处理方法和模型预测后结果的后处理方法，重新写一遍即可完成整个算法的部署迁移。

四、注意的细节

1.图像预处理尽量采用opencv直接读取图像，并对图像进行预处理。（一般情况需要将torchvisison的预处理方法改写成numpy的方法）

2.模型预测结果的后处理。（原算法中主要为torch类型的数组，而trt模型输出的结果一般为numpy类型，同样需要将处理方式从torch.tensor()修改为numpy类型）

3.模型转换完之后，模型预测时间大概在0.05s/帧，图像的前后处理也较为耗时，因此需要提升速度时，需全局考虑图像预处理操作和模型预测结果的后处理操作。