2. 模型转换

任务目标

将训练好的PyTorch分类模型转换为ONNX格式，以便后续在嵌入式平台上进行推理部署。本案例以ResNet18在ImageNet上的预训练权重为例，演示完整的转换流程、动态轴配置及正确性验证。

环境依赖

需要安装：torch、torchvision、onnx、onnxruntime。版本建议：PyTorch ≥ 1.12，ONNX ≥ 1.12，ONNX Runtime ≥ 1.14。

完整实现代码

import torch
import torchvision.models as models
import onnx
import onnxruntime as ort
import numpy as np

# 1. 加载PyTorch模型并设为评估模式
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model.eval()

# 2. 构造示例输入（batch=1, 3通道, 224x224）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
onnx_path = "resnet18.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print(f"ONNX文件已保存至: {onnx_path}")

# 4. 检查ONNX模型结构合法性
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("ONNX结构检查通过")

# 5. 用ONNX Runtime验证输出一致性
ort_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])

# 使用相同输入分别运行PyTorch和ONNX Runtime
with torch.no_grad():
    torch_output = model(dummy_input).numpy()

ort_input = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_output = ort_session.run(None, ort_input)[0]

# 6. 比较输出差异
diff = np.abs(torch_output - ort_output).max()
print(f"PyTorch与ONNX Runtime最大输出差异: {diff:.6e}")

# 7. 验证动态batch维度
dummy_batch4 = torch.randn(4, 3, 224, 224)
ort_input_b4 = {ort_session.get_inputs()[0].name: dummy_batch4.numpy()}
ort_output_b4 = ort_session.run(None, ort_input_b4)[0]
print(f"动态batch=4推理输出形状: {ort_output_b4.shape}")

# 判定标准
assert diff < 1e-4, "输出差异过大，转换可能存在问题"
assert ort_output_b4.shape[0] == 4, "动态batch维度未生效"
print("验证通过：模型转换正确，动态batch维度有效")

运行方式

将上述代码保存为 convert_onnx.py，在已安装依赖的Python环境中直接执行：

python convert_onnx.py

验证方法与通过条件

• 结构验证：onnx.checker.check_model 不抛出异常，说明ONNX图结构合法。
• 数值一致性：PyTorch与ONNX Runtime对同一输入的最大输出差异应小于 1e-4。由于浮点运算顺序差异，通常差异在 1e-6 量级，若超过 1e-3 则需检查转换配置。
• 动态轴验证：以batch=4输入ONNX Runtime，输出张量第一维应为4，证明动态batch配置生效。

程序正常结束并打印“验证通过”即表示转换成功。

常见问题

• opset版本过低：若使用旧版opset（如9），部分算子可能不支持动态轴，建议使用opset ≥ 13。
• 模型含不可导分支：若自定义模型包含条件控制流，需使用 torch.jit.trace 或 torch.onnx.export 的 trace 模式并确保输入能覆盖所有分支。
• 嵌入式部署限制：若目标平台（如某些MCU）不支持全部ONNX算子，需在转换后使用 onnxsim 简化模型，或针对目标平台进行算子适配。