前提条件:
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/帧,图像的前后处理也较为耗时,因此需要提升速度时,需全局考虑图像预处理操作和模型预测结果的后处理操作。