📚 课程目录第4章从PyTorch导出ONNX

4. 从PyTorch导出ONNX

把PyTorch模型导出为ONNX，是模型进入ONNX Runtime推理链路的第一道关口。很多“推理结果不对”“部署端加载失败”“动态batch报错”的问题，根源并不在推理引擎，而在导出阶段。本节围绕torch.onnx.export，把参数、动态轴、命名和报错四件事讲透，让你能稳定产出一份可被ONNX Runtime正确加载和优化的模型文件。

一、核心原理：导出到底做了什么

PyTorch是动态图框架，ONNX是静态计算图格式。导出的本质是：给模型喂一组示例输入（dummy input），让PyTorch执行一次前向传播，同时用追踪（Tracing）或脚本（Scripting）机制记录算子调用序列，再把这些算子翻译成ONNX标准算子，写入.onnx文件。

因此导出不是“复制模型”，而是“录制一次执行轨迹”。这意味着：控制流（if/for依赖张量值）、Python原生逻辑、未走到的分支，都不会被完整记录。理解这一点，后面所有报错和动态轴问题都能找到解释。

关键认知：导出结果 = 示例输入触发的那条执行路径 + 被翻译的算子集合。示例输入选得不好，导出的图就是错的。

二、torch.onnx.export 参数详解

函数签名较长，但真正高频使用的参数集中在下面这些。先看一个典型调用：

torch.onnx.export(
    model,                      # 待导出模型（eval模式）
    dummy_input,                # 示例输入，tuple或dict
    "model.onnx",               # 输出路径
    export_params=True,         # 是否导出权重
    opset_version=17,           # 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"}
    },
    verbose=False,
    training=torch.onnx.TrainingMode.EVAL
)

各参数的作用与取舍如下表：

参数	作用	实践建议
export_params	是否把权重写入文件	部署必须为True
opset_version	ONNX算子集版本	与推理端对齐，常用11/13/17
do_constant_folding	常量折叠，减小图规模	保持True
input_names / output_names	为输入输出命名	务必显式指定，便于后续绑定
dynamic_axes	声明哪些维度可变	动态batch/序列必填
training	导出训练图还是推理图	推理用EVAL
verbose	打印导出细节	排错时临时开启

三、动态轴：让模型支持可变维度

默认情况下，导出的ONNX会把示例输入的每个维度都固化为常量。如果示例输入是[1,3,224,224]，那么导出的模型只能接受batch=1、224×224的输入。一旦推理时传入batch=8，就会报维度不匹配。

dynamic_axes就是用来解除这种固化的。它的结构是{张量名: {维度索引: 维度名称}}。维度索引从0开始，维度名称只是标签，用于可读性。

dynamic_axes={
    "input":  {0: "batch", 2: "h", 3: "w"},
    "output": {0: "batch"}
}

动态轴只声明“可变”，不声明“范围”。ONNX本身不校验上下限，实际约束由推理端或业务逻辑保证。若需要固定范围，可在ONNX Runtime侧做输入校验。

需要注意：动态轴必须与input_names中的名字一致，否则声明无效。命名和动态轴是强耦合的，这也是下一节要强调命名的原因。

四、输入输出命名：别让下游“找不到张量”

ONNX Runtime在绑定输入、读取输出时，依赖的是张量名称。如果不显式命名，PyTorch会生成类似input.1、onnx::Gemm_123这样的自动名称，既难读又容易在版本变化时改变，导致下游代码失效。

命名要点：

输入输出名必须唯一，且与dynamic_axes中的键完全一致。
多输入模型（如BERT的input_ids、attention_mask）要按顺序命名，顺序与dummy_input的tuple顺序对应。
输出名建议语义化，如logits、boxes、scores，方便后处理。
命名后可用Netron打开确认，或通过代码读取校验。

import onnx
m = onnx.load("model.onnx")
for i in m.graph.input:
    print("输入:", i.name, [d.dim_param or d.dim_value for d in i.type.tensor_type.shape.dim])
for o in m.graph.output:
    print("输出:", o.name)

五、常见报错与复盘

导出阶段的报错大致分四类，掌握定位思路比死记报错信息更重要。

算子不支持：UnsupportedOperatorError或Exporting operator ... is not supported。原因通常是opset过低或用了自定义算子。解决：提高opset，或改写为ONNX支持的等价算子。
动态维度冲突：推理时报Got invalid dimensions for input。原因多为dynamic_axes未声明或名称不匹配。解决：核对命名与索引。
控制流丢失：导出成功但结果错误。原因是if/for依赖张量值，Tracing只记录了一条路径。解决：改用torch.jit.script或torch.onnx.export(dynamo=True)。
dtype不匹配：输入为float64而模型期望float32。解决：确保dummy_input与推理输入dtype一致。

导出前务必执行model.eval()并配合torch.no_grad()。否则Dropout、BatchNorm会以训练模式导出，导致推理结果随机或统计量错误，这是最隐蔽也最常见的坑。

复盘建议：每次导出后，用ONNX Runtime加载并跑一遍与PyTorch相同的输入，对比输出误差（通常应小于1e-4）。这一步能提前拦截绝大多数导出问题，而不是等到部署端才发现。

import onnxruntime as ort, numpy as np
sess = ort.InferenceSession("model.onnx")
out = sess.run(None, {"input": x.numpy()})
print(np.max(np.abs(out[0] - torch_out.numpy())))