pytorch模型转为tensorrt engine的通用流程

基本流程

pytorch模型->onnx模型->（onnx模型简化）->tensorrt模型

pytorch模型导出为onnx模型

dummy_input = torch.randn(50, 3, 1280, 1280, device='cuda') # 导出时推理输入
input_names = ["images"] # 输入名
output_names = ['num_dets', 'det_boxes'] # 两个输出名
save_path = '/'.join(pt_file.split('/')[:-1])
onnx_file = os.path.join(save_path, 'yolov5m.onnx')
dynamic = {'images': {0: 'batch'}} #动态batch设置
dynamic['num_dets'] = {0: 'batch'}
dynamic['det_boxes'] = {0: 'batch'}
torch.onnx.export(
    q_model,
    dummy_input,
    onnx_file,
    verbose=False,
    opset_version=13,
    do_constant_folding=False,
    input_names=input_names,
    output_names=output_names,
    dynamic_axes=dynamic) # 静态输入时，无需设置

注意事项：

• 当导出的onnx模型与原始pytorch模型测试差异较大时，大概率是因为导出模型时输入（代码中的dummy_input）引起的，将随机input变为任意一张真实数据的输入即可解决问题；

onnx模型简化

• 安装onnxsim库

  pip install onnxsim
  

• onnx模型简化

  onnxsim input_onnx.onnx output_onnx.sim
  

注意事项：

• onnx模型简化非必须步骤；
• onnx模型简化后可能变大了：正常现象，模型大小与速度无关，模型大小变大是因为原模型中含有Tile、ConstantOfShape 等算子，这些算子被消除时会产生很大的 tensor，如果不想让模型大小变大，可以使用 --no-large-tensor 参数；
• onnx模型优化后推理速度一般会更快：冗余算子比较多的模型速度提升会比较明显，冗余算子不多的模型优化后可能不会有肉眼可见的速度提升；
• onnx模型优化后对模型精度有一定影响但影响非常轻微：如博主检测模型优化后精度会有~0.3%的下降；

onnx模型转为tensorrt

trtexec --onnx=xx.onnx --saveEngine=xx.trt --minShapes=images:1x3x1280x1280 --optShapes=images:50x3x1280x1280 --maxShapes=images:100x3x1280x1280 --fp16 --workspace=4096 # 动态转换

注意事项：

• 大多数模型直接导出为fp16对模型精度几乎没有影响，但模型大小、推理访存量会降低很多，推理速度也要快很多，因此大多数模型可以直接导出为fp16精度。
• 关于tensorrt int8量化相关方法可参考tensorrt官方int8量化方法汇总