深度学习图像处理项目：PyTorch模型转ONNX格式实战解析-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00912/article/details/148361270

深度学习图像处理项目：PyTorch模型转ONNX格式实战解析

deep-learning-for-image-processing deep learning for image processing including classification and object-detection etc. 项目地址: https://gitcode.com/gh_mirrors/de/deep-learning-for-image-processing

前言

在深度学习模型部署过程中，模型格式转换是一个关键步骤。本文将详细解析如何将PyTorch训练的ResNet34模型转换为ONNX格式，这是模型优化和部署的重要前置工作。

ONNX格式简介

ONNX(Open Neural Network Exchange)是一种开放的神经网络交换格式，它允许在不同框架(PyTorch、TensorFlow等)之间转换模型。ONNX格式具有以下优势：

跨平台兼容性：可在不同推理引擎上运行
标准化：统一的模型表示格式
优化支持：便于后续进行模型优化和加速

转换流程详解

1. 准备工作

首先需要准备训练好的PyTorch模型权重文件。在本例中，使用的是针对花卉分类任务微调过的ResNet34模型，权重文件名为"resNet34(flower).pth"。

weights_path = "resNet34(flower).pth"
onnx_file_name = "resnet34.onnx"

2. 模型加载与初始化

创建ResNet34模型实例并加载预训练权重：

model = resnet34(pretrained=False, num_classes=5)
model.load_state_dict(torch.load(weights_path, map_location='cpu'))
model.eval()

关键点说明：

pretrained=False表示不加载ImageNet预训练权重
num_classes=5对应花卉数据集的5个类别
model.eval()将模型设置为评估模式，关闭dropout等训练专用层

3. 创建虚拟输入

ONNX转换需要知道模型的输入形状，这里创建一个符合预期的随机张量：

batch_size = 1
img_h = 224
img_w = 224
img_channel = 3
x = torch.rand(batch_size, img_channel, img_h, img_w, requires_grad=True)

4. 执行模型转换

使用PyTorch的torch.onnx.export函数进行转换：

torch.onnx.export(model,             # 待转换模型
                 x,                 # 模型输入
                 onnx_file_name,    # 输出文件名
                 input_names=["input"],  # 输入节点名称
                 output_names=["output"], # 输出节点名称
                 verbose=False)      # 是否打印详细信息

5. 验证转换结果

转换完成后，需要进行严格验证确保模型行为一致：

# 加载并检查ONNX模型
onnx_model = onnx.load(onnx_file_name)
onnx.checker.check_model(onnx_model)

# 使用ONNX Runtime进行推理
ort_session = onnxruntime.InferenceSession(onnx_file_name)
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)

# 对比PyTorch和ONNX Runtime输出
np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)

验证通过后会打印确认信息："Exported model has been tested with ONNXRuntime, and the result looks good!"