3、模型导出：学习如何将PyTorch模型导出为ONNX格式，包括torch.onnx.export的参数详解和常见陷阱

模型导出是ONNX边缘部署链路中的关键一环，它决定了训练框架中的模型能否被边缘推理引擎正确加载并高效执行。PyTorch作为研究界和工业界最主流的训练框架之一，其官方提供的torch.onnx.export接口承担了将动态计算图转换为静态ONNX图的重任。然而，这个转换过程并非简单的“另存为”，它涉及算子映射、动态轴声明、常量折叠、输入输出命名等一系列精细控制。如果对导出参数理解不透彻，往往会出现“训练时精度正常，导出后推理结果错乱”或“边缘设备直接报错无法加载”的窘境。本章将深入剖析torch.onnx.export的核心参数、执行逻辑以及高频陷阱，帮助读者建立一套可复用的导出检查清单。

上图展示了PyTorch模型经过torch.onnx.export转换为ONNX模型的核心流程，以及需要重点关注的参数与陷阱区域。下面我们将从导出原理、参数详解、实战案例和排障复盘四个维度展开。

3.1 导出核心原理：从动态图到静态图的“追踪”与“脚本”

PyTorch使用动态计算图（Define-by-Run），而ONNX需要静态图（Define-and-Run）。torch.onnx.export默认采用追踪（Tracing）机制：它构造一个虚拟输入，让模型真实地执行一次前向传播，同时记录下所有被调用的张量操作，并将这些操作映射为ONNX算子。这意味着，只有被实际执行到的代码路径才会被导出。如果模型内部存在依赖Python控制流（如if、for循环）且条件或循环次数取决于输入张量的值，那么追踪只会记录当前输入下的分支，导致导出的模型在边缘端遇到其他输入时行为错误。

另一种机制是脚本（Scripting），通过torch.jit.script将模型转换为TorchScript，再导出为ONNX。脚本模式能够保留控制流，但对算子支持要求更高，且部分Python特性无法使用。在边缘部署场景中，我们通常优先使用追踪模式，并通过设计模型结构（如避免数据相关的控制流）来规避其局限性。

3.2 torch.onnx.export 参数详解与配置清单

一个典型的导出调用如下所示：

import torch
import torch.onnx

# 假设 model 是训练好的 PyTorch 模型，已设置为 eval 模式
model.eval()

# 构造一个与训练时输入 shape 一致的虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出
torch.onnx.export(
    model,                      # 要导出的模型
    dummy_input,                # 虚拟输入，用于追踪计算图
    "model.onnx",               # 输出文件路径
    export_params=True,         # 是否导出模型权重，训练好的模型必须为 True
    opset_version=12,           # ONNX 算子集版本，边缘端引擎需支持
    do_constant_folding=True,   # 是否执行常量折叠优化
    input_names=["input"],      # 输入张量名称
    output_names=["output"],    # 输出张量名称
    dynamic_axes={
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "output": {0: "batch_size"}
    }                           # 声明动态轴
)

下表对关键参数进行了详细说明：

参数名	默认值	作用与推荐设置
model	必填	必须是 nn.Module 实例，且导出前应调用 model.eval() 关闭 Dropout 和 BatchNorm 的训练行为。
args	必填	可以是单个张量或元组。必须与模型前向传播的输入结构一致，且数据类型、维度需与训练时一致。
f	必填	输出文件路径或文件对象。建议使用 .onnx 后缀。
export_params	True	若为 False，则导出不包含权重的“空壳”模型，仅用于结构验证。部署时需设为 True。
opset_version	9（较老版本）	边缘端推理引擎（如 ONNX Runtime、TensorRT、OpenVINO）对 opset 有最低要求。建议使用 11-15 之间的版本，兼顾算子覆盖与兼容性。过高版本可能导致旧引擎无法加载。
do_constant_folding	True	将常量计算在导出时预先算好，减小模型体积并提升推理速度。但若模型中有 BatchNorm，折叠后可能改变其行为，需谨慎。
input_names / output_names	None	为输入输出张量指定名称。边缘端推理时通过名称绑定数据，建议使用语义化名称，如 "input"、"output"。
dynamic_axes	None	以字典形式声明哪些维度是动态的。例如 {"input": {0: "batch"}} 表示第 0 维可变。若不声明，所有维度均被固定为导出时的具体数值。
verbose	False	设为 True 可打印导出过程中的详细日志，便于排查算子映射问题。

3.3 实战案例：动态尺寸分类模型的导出与验证

假设我们有一个轻量级图像分类模型，输入为 RGB 图像，但希望边缘端能够处理不同分辨率的输入（例如 224x224 和 320x240）。如果直接使用固定尺寸导出，边缘端将无法接受其他分辨率的输入。以下是完整的导出与验证流程：

1. 准备模型与虚拟输入：加载训练好的 ResNet-18 模型，设置为 eval 模式。构造一个 1x3x224x224 的虚拟输入。
2. 声明动态轴：在 dynamic_axes 中声明第 0 维（batch）以及第 2、3 维（高、宽）为动态。
3. 执行导出：设置 opset_version=12，do_constant_folding=True，输出文件为 resnet18_dynamic.onnx。
4. 使用 ONNX Runtime 验证：编写 Python 脚本，分别用 1x3x224x224 和 2x3x320x240 的随机输入进行推理，检查输出 shape 是否符合预期。

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("resnet18_dynamic.onnx")
input_name = sess.get_inputs()[0].name

# 测试不同 shape
for shape in [(1,3,224,224), (2,3,320,240)]:
    dummy = np.random.randn(*shape).astype(np.float32)
    outputs = sess.run(None, {input_name: dummy})
    print(f"Input shape: {shape}, Output shape: {outputs[0].shape}")

运行结果应显示输出 shape 分别为 (1, 1000) 和 (2, 1000)，证明动态轴生效。若未声明动态轴，第二次推理将直接报错，提示输入维度不匹配。

3.4 常见陷阱与排障复盘

在实际导出过程中，以下四个陷阱出现频率最高，且往往导致部署失败或推理精度异常。

陷阱一：模型处于训练模式导致 BatchNorm 统计量错乱

如果导出前忘记调用 model.eval()，BatchNorm 层会使用当前 batch 的均值和方差进行归一化，而非使用训练阶段累积的全局统计量。这会导致导出的 ONNX 模型在边缘端推理时，输出结果与训练时的验证结果严重偏离。此外，Dropout 层在训练模式下会随机丢弃神经元，导致导出时追踪到的计算图包含随机性，但 ONNX 中无法表达这种随机行为，最终导出的模型行为不可控。

判断方法：导出前打印模型状态：print(model.training)，确保输出为 False。

陷阱二：动态轴声明不完整或过度声明

有些模型内部存在多个张量运算，例如在 forward 中先对输入做了 view 或 reshape，再经过全连接层。如果只声明了输入的 batch 维为动态，但全连接层要求固定输入特征维度，那么当 batch 变化时，view 操作可能计算出错误的 shape。反之，如果过度声明了某些不应变化的维度（如通道数），可能导致边缘端引擎生成低效的算子。

判断方法：导出后使用 onnx.shape_inference 工具检查各中间张量的 shape 标注，确认动态维度传播是否符合预期。

陷阱三：算子不兼容导致导出失败或推理报错

PyTorch 中部分算子（如 torch.where 的某些用法、torch.topk、自定义 CUDA 算子）在 ONNX 中没有直接对应实现。此时 torch.onnx.export 可能抛出 UnsupportedOperatorError，或者导出成功但边缘端引擎加载时报错。例如，torch.nn.functional.interpolate 的 align_corners=False 模式在 opset 11 以下不支持。

解决方案：升级 opset_version；或者使用 torch.onnx.register_custom_op_symbolic 为自定义算子注册映射函数；最直接的办法是修改模型实现，使用 ONNX 支持的等价算子组合。

陷阱四：常量折叠改变模型行为

do_constant_folding=True 会将所有输入无关的常量计算提前算好。对于包含 BatchNorm 的模型，折叠后 BatchNorm 的 scale 和 bias 会被融合到前一层的权重中，这本身是优化。但如果模型在推理时依赖 BatchNorm 的动态行为（例如用于风格迁移的 Adaptive Instance Normalization），折叠将导致错误。

判断方法：导出后对比原始 PyTorch 模型与 ONNX 模型在相同输入下的输出差异。若差异超过 1e-4，尝试设置 do_constant_folding=False 重新导出。

3.5 导出验证脚本模板

为了确保每次导出都可靠，建议将以下验证步骤固化为脚本：

import onnx
import onnxruntime as ort
import numpy as np
import torch

def verify_onnx(model, onnx_path, dummy_input, rtol=1e-3, atol=1e-5):
    # 1. 结构检查
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)

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

    # 3. ONNX Runtime 推理
    sess = ort.InferenceSession(onnx_path)
    onnx_output = sess.run(None, {sess.get_inputs()[0].name: dummy_input.numpy()})[0]

    # 4. 数值比对
    np.testing.assert_allclose(torch_output, onnx_output, rtol=rtol, atol=atol)
    print("Verification passed: PyTorch and ONNX outputs match.")

该脚本应作为 CI/CD 流水线中的一环，在每次模型更新后自动执行，防止回归问题。

3.6 总结复盘

模型导出是连接训练与部署的桥梁，其核心在于理解追踪机制的本质——记录实际执行路径。通过合理配置 dynamic_axes、opset_version 和 do_constant_folding，并配合严格的验证脚本，可以规避绝大多数导出问题。建议读者在项目初期就建立导出检查清单，包括：模型是否处于 eval 模式、输入输出命名是否清晰、动态轴是否与边缘端需求匹配、算子兼容性是否验证、数值一致性是否通过。将这些步骤固化到自动化流程中，能够显著降低边缘部署的试错成本。