4、导出前的模型检查：模型结构验证、输入输出定义、动态轴设置

在正式调用 torch.onnx.export() 或 onnx.save() 之前，有一道极易被忽视却至关重要的工序——导出前的模型检查。许多开发者习惯将代码写好、跑通一次前向推理后就直接导出，结果在 ONNX Runtime 或 TensorRT 中遇到各种“莫名其妙”的报错。事实上，绝大多数导出报错并非源于导出过程本身，而是模型在“出发”前就没有做好体检。本节将围绕三个核心检查维度展开：模型结构验证、输入输出定义、动态轴设置，帮助你建立一套可复用的导出前检查清单。

4.1 核心原理：导出不是“另存为”，而是“重新编译”

理解导出前检查的必要性，首先要纠正一个认知偏差：ONNX 导出并不是像 Word 文档“另存为 PDF”那样简单的格式转换。PyTorch 模型导出为 ONNX 时，实际上经历了一次计算图追踪（tracing）或脚本化（scripting）的过程。追踪模式会实际执行一遍模型，将实际运行到的算子记录下来，转换为 ONNX 的节点图。这意味着：

• 数据相关的控制流会被固化：例如 if tensor.sum() > 0 这类依赖张量数值的分支，在追踪时只会保留当前执行路径，另一条分支会被永久丢弃。
• 动态 shape 可能被静态化：如果输入张量的某个维度在追踪时是固定值（如 batch=8），ONNX 图中该维度可能被写死为 8，后续推理时若传入 batch=1 就会报错。
• Python 原生操作无法直接转换：如 list.append、dict 遍历、numpy 随机操作等，需要被 PyTorch 算子替代或通过 torch.onnx.is_onnx_supported() 检查。

因此，导出前的模型结构验证，本质上是在回答三个问题：图是否完整？算子是否可转换？shape 是否明确？ 下面我们逐一展开。

4.2 模型结构验证：从“能跑”到“能导出”

很多模型在 PyTorch 中能正常前向推理，但导出时却报“Unsupported operator”或“Constant folding failed”。这是因为 PyTorch 的 eager 模式非常宽容，允许 Python 的灵活性渗透到计算图中；而 ONNX 是一个静态图格式，要求所有操作都能映射到 ONNX 算子集。

结构验证建议按以下步骤执行：

1. 检查模型类型：确认待导出对象是 nn.Module 子类实例，且已调用 model.eval()。训练模式下的 BatchNorm、Dropout 等层会引入随机性或依赖统计量，导致导出图与推理行为不一致。
2. 遍历子模块：使用 model.named_modules() 检查是否有自定义层或包含 Python 控制流的模块。对于自定义 autograd.Function，需要额外实现 symbolic 方法。
3. 算子兼容性预检：PyTorch 1.10+ 提供了 torch.onnx.export(..., dynamo=True) 或使用 onnx.checker 工具。更直接的方式是设置 torch.onnx.export(..., verbose=True)，观察日志中是否有 “Warning: Unsupported operator” 字样。
4. 使用 onnx.checker 验证：导出后立即调用 onnx.checker.check_model(model_proto)，它会检查图的结构完整性、节点输入输出是否匹配、属性是否合法等。

下面是一个典型的验证代码片段：

import torch
import onnx

model = MyModel().eval()
dummy_input = torch.randn(1, 3, 224, 224)

# 导出时开启 verbose 和 dynamo（如可用）
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=17,
    dynamo=True,          # 若 PyTorch ≥ 2.0
    dynamic_axes=None,    # 本节后续会讲
    verbose=True
)

# 结构完整性检查
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 模型结构验证通过")

4.3 输入输出定义：名称、顺序与数据类型

ONNX 模型是一个完全独立的计算图，它不依赖 PyTorch 的 Python 对象。因此，输入输出的名称、顺序和数据类型必须显式声明，否则下游推理引擎（如 ONNX Runtime）将无法正确绑定数据。

输入输出定义的关键点包括：

• input_names / output_names：建议使用语义化名称，如 "input"、"output" 或 "pixel_values"、"logits"。名称一旦确定，后续推理时需严格匹配。
• 数据类型：PyTorch 中默认的 torch.float32 会映射为 ONNX 的 tensor(float)。若模型包含 torch.int64 的输入（如 token id），需确保 dummy input 的类型正确，否则导出图可能默认使用 float。
• 输入顺序：如果模型有多个输入，dummy_input 可以是一个元组，顺序必须与 input_names 一一对应。例如：

dummy_input = (torch.randn(1, 3, 224, 224), torch.tensor([1], dtype=torch.long))
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["image", "label"],
    output_names=["logits"],
    opset_version=17
)

此外，建议在导出后使用 onnxruntime 进行一次推理验证，确认输入输出的 shape 和 dtype 与预期一致：

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("model.onnx")
for inp in sess.get_inputs():
    print(f"输入名: {inp.name}, shape: {inp.shape}, 类型: {inp.type}")
for out in sess.get_outputs():
    print(f"输出名: {out.name}, shape: {out.shape}, 类型: {out.type}")

4.4 动态轴设置：让模型不再“死板”

默认情况下，ONNX 导出会将所有输入输出的维度视为静态值。如果训练时 batch size 固定为 8，导出后模型就只能接受 batch=8 的输入。但在实际部署中，推理请求的 batch 往往是动态变化的（如在线服务一次只处理 1 张图）。此时必须通过 dynamic_axes 参数声明哪些维度是动态的。

动态轴设置的核心语法如下：

dynamic_axes = {
    "input": {0: "batch_size", 2: "height", 3: "width"},
    "output": {0: "batch_size"}
}

torch.onnx.export(
    model,
    dummy_input,
    "model_dynamic.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes=dynamic_axes,
    opset_version=17
)

上述配置的含义是：input 的第 0 维（batch）、第 2 维（height）、第 3 维（width）是动态的，分别命名为 batch_size、height、width；output 的第 0 维也是动态的，且与输入共享 batch_size 这个符号名称。这样，ONNX Runtime 在推理时就能接受任意 batch 和任意 H/W 的输入（只要模型本身支持）。

下表总结了常见动态轴设置场景及建议：

模型类型	动态维度	推荐设置	注意事项
CNN 图像分类	batch, H, W	batch 设为动态，H/W 视部署需求	若模型含 AdaptiveAvgPool，H/W 可动态
NLP Transformer	batch, sequence_length	两者均设为动态	attention mask 需同步动态
目标检测（两阶段）	batch, 候选框数量	batch 动态，候选框数量通常动态	后处理部分可能需自定义算子
语音识别	batch, 时间帧数	时间帧设为动态	需确保模型内部无固定长度假设

4.5 案例：一个典型的导出前检查流程

假设我们有一个基于 ResNet 的图像分类模型，期望部署到支持动态 batch 的推理服务中。完整的导出前检查流程如下：

1. 结构验证：打印模型结构，确认没有自定义算子或 Python 控制流。
2. 输入输出定义：明确输入为 images（float32, NCHW），输出为 logits（float32, NC）。
3. 动态轴设置：将 batch 维设为动态，H/W 保持静态（因为 ResNet 的输入尺寸固定为 224×224）。
4. 导出并验证：使用 onnx.checker 和 onnxruntime 分别验证结构正确性和推理一致性。

以下是最终导出代码：

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

model = torchvision.models.resnet18(pretrained=True).eval()
dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18_dynamic.onnx",
    input_names=["images"],
    output_names=["logits"],
    dynamic_axes={"images": {0: "batch"}, "logits": {0: "batch"}},
    opset_version=17
)

# 结构验证
onnx_model = onnx.load("resnet18_dynamic.onnx")
onnx.checker.check_model(onnx_model)

# 动态 batch 推理验证
ort_sess = ort.InferenceSession("resnet18_dynamic.onnx")
for batch in [1, 4, 8]:
    test_input = np.random.randn(batch, 3, 224, 224).astype(np.float32)
    outputs = ort_sess.run(None, {"images": test_input})
    assert outputs[0].shape[0] == batch
    print(f"batch={batch} 验证通过，输出 shape: {outputs[0].shape}")

4.6 误区与复盘：导出前检查的五个“坑”

在实际教学和项目支持中，以下五个问题出现频率最高，值得反复强调：

• 坑 1：忘记调用 model.eval()。训练模式下 BatchNorm 会使用 batch 统计量，导致导出的图与推理行为不一致。导出前务必切换为 eval 模式。
• 坑 2：dummy input 的 shape 过于随意。如果 dummy input 的某个维度为 1，而该维度在后续推理中可能变化，必须显式声明为动态轴，否则导出图会将该维度固定为 1。
• 坑 3：忽略 opset_version 的影响。不同算子集版本支持的算子不同，较新的算子（如 GridSample）需要 opset ≥ 16。建议使用 opset 17 或更高版本，但需确认部署环境支持。
• 坑 4：动态轴名称不一致。如果输入和输出的 batch 维都声明为动态，但名称不同（如一个叫 batch，另一个叫 batch_size），ONNX Runtime 会认为它们是独立的动态维度，可能导致 shape 推断失败。
• 坑 5：只验证导出成功，不验证推理一致性。导出成功只代表图结构合法，不代表数值正确。务必用同一输入分别跑 PyTorch 和 ONNX Runtime，比较输出差异是否在容差范围内（如 np.allclose(atol=1e-5)）。