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

最新推荐文章于 2024-09-22 16:59:42 发布

要坚持写博客呀

最新推荐文章于 2024-09-22 16:59:42 发布

阅读量1k

点赞数

文章标签： pytorch 深度学习 python

本文链接：https://blog.csdn.net/weixin_39263657/article/details/130606346

版权

本文介绍了如何将PyTorch的ResNet50模型转换为ONNX格式，包括设置BatchSize，使用torch.onnx.export进行转换，验证ONNX模型格式，简化模型以及对比转换前后的精度。此外，还提到了使用ONNXRuntime进行模型验证和NETRON进行模型可视化。

摘要由CSDN通过智能技术生成

一、转换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!')