1. PyTorch模型转换为ONNX基础

任务目标

将一个简单的PyTorch CNN分类模型转换为ONNX格式，并通过对比转换前后模型在相同输入下的输出，验证转换正确性。

环境依赖

• Python 3.8+
• PyTorch 2.0+（含 torch.onnx）
• ONNX Runtime（用于加载和推理ONNX模型）

安装命令：pip install torch onnxruntime

关键完整代码

import torch
import torch.nn as nn
import onnxruntime as ort
import numpy as np

# 1. 定义简单CNN模型
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        return self.classifier(self.features(x))

# 2. 初始化模型并设为推理模式
model = SimpleCNN()
model.eval()

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

# 4. 转换为ONNX格式
onnx_path = "simple_cnn.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},  # 支持动态batch
    opset_version=17,
)

# 5. 使用ONNX Runtime加载并推理
ort_session = ort.InferenceSession(onnx_path)
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)[0]

# 6. 对比PyTorch与ONNX输出
with torch.no_grad():
    torch_output = model(dummy_input).numpy()

# 7. 验证一致性
max_diff = np.max(np.abs(torch_output - ort_outputs))
print(f"PyTorch输出形状: {torch_output.shape}")
print(f"ONNX输出形状: {ort_outputs.shape}")
print(f"最大误差: {max_diff:.6f}")

# 判定标准：误差小于1e-5即通过
assert max_diff < 1e-5, "转换验证失败：输出差异过大"
print("✓ 转换验证通过：PyTorch与ONNX输出一致")

运行与验证

1. 将上述代码保存为 convert_to_onnx.py，直接运行：python convert_to_onnx.py
2. 程序执行后会在当前目录生成 simple_cnn.onnx 文件。
3. 控制台输出包含两个输出张量的形状以及最大误差值。

通过条件

• 程序无异常退出，且打印 ✓ 转换验证通过。
• 最大误差小于 1e-5，说明ONNX模型与原始PyTorch模型在数值上保持一致。
• 生成的 simple_cnn.onnx 文件可被ONNX Runtime正常加载（脚本第5步已隐式验证）。

常见问题

• 动态轴警告：若不需要动态batch，可删除 dynamic_axes 参数，模型将固定为batch=1。
• opset版本过低：若环境中ONNX Runtime版本较旧，可将 opset_version 降至11或13。
• 误差过大：检查模型是否包含训练/推理模式差异（如Dropout、BatchNorm），确保已调用 model.eval()。