2、PyTorch导出ONNX核心API：torch.onnx.export()参数详解、动态轴与静态轴设置

好，咱们直接进入正题。这一章我带你死磕 torch.onnx.export() 这个函数。说实话，很多同学模型写得很溜，一到导出就翻车，十有八九是没搞懂这个API的脾气。

我个人习惯把导出ONNX比作「给模型拍证件照」。你模型长什么样，输入输出是什么尺寸，都得交代清楚。拍歪了、尺寸不对，后面部署推理的时候全得重来。

2.1 核心参数逐个拆解

先看最常用的签名，我直接给你列出来：

torch.onnx.export(
    model,               # 要导出的模型
    args,                # 模型的输入张量（或元组）
    f,                   # 输出文件路径，比如 "model.onnx"
    export_params=True,  # 是否导出模型权重
    verbose=False,       # 是否打印导出日志
    input_names=None,    # 输入节点的名字
    output_names=None,   # 输出节点的名字
    dynamic_axes=None,   # 动态轴配置（重点！）
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True  # 是否做常量折叠优化
)

嗯，这里要注意，参数看着多，但真正让你头疼的就那么几个。我一个个说。

model 和 args

model 就是你的 PyTorch 模型实例，记得要调成 eval() 模式。为什么？因为 dropout 和 batchnorm 在训练和推理时的行为不一样。我曾经见过有人忘了切模式，导出的 ONNX 推理结果跟 PyTorch 对不上，排查了半天。

args 是模型的输入。可以是单个张量，也可以是元组。比如你的模型有两个输入，就传 (input1, input2)。这里有个坑：传入的输入数据维度要跟你实际部署时一致。你想想看，如果你传了个 batch=1 的输入去 trace，那导出的模型就默认 batch 是 1，后面想改都难。

input_names 和 output_names

这两个参数给计算图的输入输出节点起名字。名字起得好，后面推理引擎调用时就清爽。我一般习惯用 "input"、"output"，或者更语义化的 "image"、"boxes"。

小技巧：名字不要用中文，也不要用特殊符号。ONNX 规范里只支持字母、数字和下划线。我之前有个同事用了横杠，结果在 TensorRT 里死活解析不出来。

opset_version

这个参数决定了 ONNX 算子集的版本。版本越高，支持的算子越新，但兼容性可能差一些。我个人建议：能用 11 就用 11，这是最稳妥的版本。如果你用了 Transformer 或者一些很新的操作，可能需要升到 13 或 14。但记住，版本越高，对推理引擎的要求也越高。

2.2 动态轴 vs 静态轴：核心中的核心

好，接下来是本章的重头戏。说白了，静态轴就是固定死的尺寸，动态轴就是可以变的尺寸。

举个例子。你训练时 batch size 是 32，导出 ONNX 时如果不做任何设置，那导出的模型就只接受 batch=32 的输入。部署时来了一个样本，你还得凑够 32 个才能推理，这不扯吗？

所以我们需要动态轴。最常见的场景就是 batch size 动态 和 序列长度动态。

怎么设置动态轴？

看代码：

# 假设模型输入是 [batch, 3, 224, 224]
# 输出是 [batch, 1000]

dynamic_axes = {
    "input": {0: "batch_size"},   # 输入的第0维是动态的，叫 batch_size
    "output": {0: "batch_size"}  # 输出的第0维也动态，保持一致
}

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

这里 dynamic_axes 是一个字典。key 是节点的名字（就是上面 input_names 和 output_names 里定义的名字），value 又是一个字典，key 是维度索引，value 是给这个维度起的名字。

为什么要起名字？因为 ONNX 内部需要知道哪些维度是联动的。比如输入 batch 变了，输出 batch 也得跟着变，它们都叫 batch_size，ONNX 就知道这两个维度是一回事。

多维度动态

有时候不止 batch 要动态。比如 NLP 模型，序列长度也是动态的：

dynamic_axes = {
    "input_ids": {0: "batch_size", 1: "seq_len"},
    "attention_mask": {0: "batch_size", 1: "seq_len"},
    "output": {0: "batch_size", 1: "seq_len"}
}

这里我把输入和输出的第0维和第1维都设成了动态。注意，同一个维度名字在不同节点之间要对应好。比如 seq_len 在输入和输出里都出现，ONNX 就知道它们应该保持一致。

警告：不是所有算子都支持动态轴。比如 Reshape 操作，如果你把某个维度设成动态，但后面跟着一个固定形状的 Reshape，那导出时可能会报错。我曾经在导出一个多尺度检测模型时，就因为一个 Flatten 操作卡了两天。最后发现是动态轴和静态 reshape 冲突了。

2.3 静态轴的适用场景

是不是所有场景都要用动态轴？当然不是。静态轴也有它的优势：性能更好。

你想想看，推理引擎拿到一个静态形状的模型，可以做很多编译优化，比如内存预分配、算子融合。而动态形状的模型，每次推理都要重新计算一些东西，性能会打折扣。

所以我的建议是：

  • 生产环境、性能敏感：尽量用静态轴。比如你的服务只接受固定尺寸的图片，那就别折腾动态了。
  • 需要灵活部署：用动态轴。比如你的模型要处理不同长度的文本。
  • 实在拿不准：先静态导出，跑通了再改动态。一步步来，别一口吃成胖子。

2.4 实战中的常见坑

我总结几个我踩过的坑，你遇到了可以直接翻到这里查：

问题 原因 解决方案
导出时报错 "RuntimeError: Only tuples, lists and Variables are supported..." 输入类型不对，传了 dict 或者自定义类 把输入转成 torch.Tensor 或 tuple of Tensor
导出的 ONNX 推理结果和 PyTorch 不一致 模型没切 eval 模式，或者有动态控制流 加 model.eval()，用 torch.jit.trace 替代 torch.onnx.export
动态轴设置后推理报错 "Shape mismatch" 某个中间算子不支持动态形状 用 Netron 可视化 ONNX 图，找到不支持动态的节点，手动替换
导出的 ONNX 文件巨大 模型里用了大量常数或者循环展开 开启 do_constant_folding=True，或者减小 opset_version
核心总结：torch.onnx.export() 的核心就三件事——告诉它模型是什么、输入长什么样、哪些维度可以变。动态轴不是银弹，静态轴也不是万能的。根据你的部署场景，选最合适的策略。

嗯，这一章的内容就到这。下一章我会带你手把手导出一个实际模型，从报错到解决，把今天讲的参数全部实战一遍。到时候你就知道，纸上得来终觉浅，绝知此事要躬行。

← 上一章 📖 返回目录 下一章 →