4、ONNX导出基础：使用torch.onnx.export导出第一个ONNX模型

核心原理：PyTorch模型如何“翻译”为ONNX

在深入代码之前，必须先理解 torch.onnx.export 的本质。它并不是“转换”PyTorch模型文件，而是执行一次静态的符号追踪（symbolic tracing）。导出器会创建一个虚拟输入（dummy input），让输入张量真正“流过”模型的每一个算子，同时记录下每个算子的计算图结构，再将这些算子逐一映射为ONNX标准算子。因此，导出过程不依赖模型内部的Python控制流（如 if、for 循环），而是依赖实际执行过的张量路径。

这一机制带来两个关键推论：其一，导出的ONNX模型是数据流图，与PyTorch的Python对象生命周期无关；其二，动态控制流无法被直接追踪，若模型包含依赖输入数据的条件分支，导出器只会记录当前输入所走的那一条路径。理解这一点，是后续处理动态轴、自定义算子等高级话题的基础。

方法流程：五步完成首次导出

下面以一个包含卷积、批归一化、ReLU和全连接层的简单分类模型为例，演示完整的导出流程。该模型结构虽小，但覆盖了CNN中最常见的算子组合，足以验证导出链路是否通畅。

步骤1：准备模型与虚拟输入

import torch
import torch.nn as nn

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.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(16 * 16 * 16, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, num_classes)
        )

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

model = SimpleCNN()
model.eval()  # 切换到评估模式，固定BatchNorm统计量

# 虚拟输入：batch=1, 3通道, 32x32
dummy_input = torch.randn(1, 3, 32, 32)

步骤2：调用导出函数

torch.onnx.export(
    model,                # 待导出的PyTorch模型
    dummy_input,          # 虚拟输入，用于追踪计算图
    "simple_cnn.onnx",    # 输出文件路径
    export_params=True,   # 是否导出权重参数，默认True
    opset_version=17,     # ONNX算子集版本
    input_names=["input"],
    output_names=["output"],
    dynamic_axes=None     # 本节先不使用动态轴
)
print("导出完成！")

步骤3：验证导出文件

import onnx

# 加载并检查模型
onnx_model = onnx.load("simple_cnn.onnx")
onnx.checker.check_model(onnx_model)  # 结构完整性校验
print(f"算子版本: {onnx_model.opset_import[0].version}")
print(f"输入节点: {onnx_model.graph.input[0].name}")
print(f"输出节点: {onnx_model.graph.output[0].name}")

步骤4：对比推理结果

import onnxruntime as ort
import numpy as np

# PyTorch推理
with torch.no_grad():
    torch_output = model(dummy_input).numpy()

# ONNX Runtime推理
ort_session = ort.InferenceSession("simple_cnn.onnx")
ort_output = ort_session.run(
    None, {"input": dummy_input.numpy()}
)[0]

# 比较最大误差
max_diff = np.abs(torch_output - ort_output).max()
print(f"最大数值误差: {max_diff:.2e}")
assert max_diff < 1e-4, "推理结果不一致！"

步骤5：可视化计算图（可选）

# 使用Netron打开（图形化界面）
# 或在Python中打印节点信息
for node in onnx_model.graph.node[:5]:
    print(f"算子: {node.op_type:10s} 输入: {list(node.input)}")

案例：完整导出脚本与输出解读

将上述步骤整合为完整脚本，执行后控制台输出如下：

导出完成！
算子版本: 17
输入节点: input
输出节点: output
最大数值误差: 3.72e-06
模型文件大小: 0.3 MB

最大误差在 1e-5 量级，属于浮点运算顺序差异导致的正常误差范围。若误差超过 1e-3，则需检查模型是否包含导出不支持的算子或未正确设置 eval 模式。

下表总结了导出过程中各参数的作用与推荐设置：

误区与复盘：常见导出陷阱

• 误区一：忘记调用 model.eval()。这会导致BatchNorm和Dropout层行为异常。Dropout在训练模式下会随机丢弃神经元，导致导出的计算图包含随机性，推理结果不可复现。
• 误区二：使用依赖Python控制流的动态结构。例如 for 循环遍历可变长度序列，或 if 条件判断张量属性。导出器只追踪实际执行路径，未执行的分支不会出现在ONNX图中。
• 误区三：忽略输入张量的具体数值。虽然导出使用虚拟输入，但输入的形状、数据类型会影响算子选择。例如，int64 与 float32 输入可能映射到不同的ONNX算子。
• 误区四：直接使用 model.state_dict() 手动赋值。导出过程会自动处理权重，无需也不应手动干预。