TensorRT推理——（一）转换成onnx模型

一、 转换ONNX步骤

使用TensorRT推理肯定是在英伟达GPU环境上进行推理的。

• 设置Batch Size
• 转换模型
• 验证转换前后精度

二、转换成onnx

1. torch.onnx.export转换ONNX

pytoch计算图是静态的，先编译后执行。

import torch
import torch.onnx
import torchvision.models as models
# 输入sahpe (和输入shape一样，数据内容随机，Tensor类型)
dummy_input=torch.randn(BATCH_SIZE, 3, 224, 224)
# device设置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dummy_input = dummy_input.to(device)
# 加载model
resnet50_model = models.resnet50(pretrained=True, progress=False).to("cuda").eval()
# 设置Batch Size
BATCH_SIZE = 32

# 使用torch.onnx.export转换
torch.onnx.export(resnet50_model,           # loaded torch model
                  dummy_input,              # input shape
                  "resnet50_pytorch.onnx",  # export onnx model name path
                  opset_version=11,         # op version
                  training=torch.onnx.TrainingMode.EVAL,  # 以推理模式导出模型
                  do_constant_folding=True, # 是否使用“常量折叠”优化。常量折叠将使用一些算好的常量来优化一些输入全为常量的节点
                  input_names=["input"],    # 按顺序分配onnx图输入节点名称，可有多个输入
                  output_names=["output"],  # 按顺序分配onnx图输出节点名称，可有多个输出
                  # dynamic_axes: 设置动态维度，key为input_names output_names的名字，value中的key为对应哪些维度是动态的，value中的value为给对应维度取的名称
                  dynamic_axes={"input": {0:"batch_size", 1:"channel", 2:"height", 3:"height"}},
                  verbose=False)   

2. 验证导出的ONNX模型格式是否正确

如果有错误的话该函数会直接报错。模型正确的话，控制台中应该会打印出"Model correct"。

onnx_model = onnx.load(onnx_model_path)  # load onnx model
onnx.checker.check_model(onnx_model)  # check onnx model

使用onnxsim对onnx模型后进行简化。就是把conv和bn合在一块，以及一些计算出来是常量的模块直接折叠之类的，可以加快onnx模型推理。如果你需要的时tensorrt，那么可能不太需要这步，因为tensorrt也做了这些优化。

import onnxsim
onnx_model = onnx.load(ONNX_MODEL_PATH)
onnx_sim_model, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(onnx_sim_model, ONNX_SIM_MODEL_PATH)
print('ONNX file simplified!')

对动态batch的onnx模型进行简化:

import onnxsim
onnx_model = onnx.load(ONNX_MODEL_PATH)
onnx_sim_model, check = simplify(onnx_model, dynamic_input_shape=True, input_shapes={"images":[1, 3, 224, 224]})
assert check, "Simplified ONNX model could not be validated"
onnx.save(onnx_sim_model, ONNX_SIM_MODEL_PATH)
print('ONNX file simplified!')

3. 验证模型转换前后的精度

使用onnxruntime和onnxruntime-gpu跑一下模型，前、后处理要和转换前的模型一致。
也有将EfficientNMS作为插件做成一个End to End的模型进行推理。

4. 将转换前后的模型放入到NETRON上查看转换前后的差别，便于排错。

ref:

https://blog.csdn.net/m0_38043555/article/details/114479282
https://github.com/Linaom1214/TensorRT-For-YOLO-Series