PyTorch模型导出ONNX：核心参数与动态轴设置

说实话，很多同学在模型部署时卡住，不是因为模型本身多复杂，而是卡在导出这一步。torch.onnx.export 这个函数，参数看着不多，但每个都藏着坑。今天我就把几个核心参数掰开揉碎了讲，顺便聊聊动态轴和静态轴到底怎么选。

一、torch.onnx.export 核心参数详解

先看一个最基础的导出代码，我习惯这样写：

import torch
import torch.onnx

# 定义模型
model = MyModel()
model.eval()

# 构造一个 dummy input
dummy_input = torch.randn(1, 3, 224, 224)

# 导出
torch.onnx.export(
    model,                     # 模型
    dummy_input,               # 示例输入
    "model.onnx",              # 输出路径
    export_params=True,        # 是否导出参数
    opset_version=11,          # opset 版本
    do_constant_folding=True,  # 常量折叠
    input_names=["input"],     # 输入名称
    output_names=["output"],   # 输出名称
    dynamic_axes={}            # 动态轴设置
)

这里每个参数都有讲究，我挑几个容易出问题的重点说。

1. export_params：要不要带权重？

默认是 True，导出模型权重。如果你只想导出网络结构做可视化，可以设成 False。但部署时一定要 True，否则推理结果全是错的。我曾经帮一个同事排查问题，他导出的 ONNX 跑出来全是随机数，折腾了半天才发现是 export_params=False。

2. opset_version：版本选不对，算子就炸

ONNX 每个版本支持的算子集合不一样。我个人的建议：

  • 一般场景选 opset=11，兼容性好，大部分算子都支持
  • 如果用了较新的 PyTorch 算子（比如一些变形操作），可能需要 opset=13 或更高
  • 不要盲目追新，opset 越高，对推理引擎的版本要求也越高
小技巧：如果你不确定选哪个版本，先试 opset=11。导出报错时看日志，它会告诉你哪个算子不支持，再考虑升级 opset。

3. do_constant_folding：剪枝优化，但小心副作用

这个参数默认 True，会把一些计算图里可以提前算好的常量折叠掉。比如模型里有个固定的 reshape 操作，折叠后 ONNX 里就直接变成常量了。

但要注意：如果你的模型里有条件分支或者动态形状相关的操作，常量折叠可能会把逻辑写死。我遇到过一个问题：模型里有个根据输入尺寸动态调整的 padding 操作，常量折叠后 padding 值被固定了，换不同尺寸的输入就报错。所以如果你的模型对输入形状敏感，建议先关掉这个参数试试。

4. input_names / output_names：命名规范很重要

这两个参数看似简单，但命名不规范会坑到后续的推理引擎。我建议：

  • 输入名称用有意义的词，比如 "input", "image", "data"
  • 输出名称用 "output", "logits", "boxes" 等
  • 不要用特殊字符或中文，有些推理引擎解析会出问题

二、动态轴与静态轴：到底怎么选？

这是导出 ONNX 时最让人纠结的问题。说白了：

  • 静态轴：输入输出的形状在导出时就固定了，推理时不能变
  • 动态轴：某些维度可以在推理时变化，比如 batch size 或序列长度

先看一个对比表格：

对比项 静态轴 动态轴
推理速度 更快（算子可做形状优化） 稍慢（需要处理动态形状）
灵活性 差，输入形状必须固定 好，支持不同尺寸输入
部署难度 低，推理引擎容易优化 高，部分引擎支持不完善
适用场景 生产环境、固定 batch 推理 需要变 batch、变分辨率、变序列长度

静态轴的设置

默认情况下，torch.onnx.export 导出的就是静态轴。你 dummy_input 是什么形状，导出的 ONNX 就固定成什么形状。比如：

dummy_input = torch.randn(1, 3, 224, 224)
# 导出的 ONNX 只接受 [1, 3, 224, 224] 的输入

如果你尝试传入 [2, 3, 224, 224] 的输入，推理引擎会直接报错。嗯，这就是静态轴的特点——简单、高效、但死板。

动态轴的设置

动态轴通过 dynamic_axes 参数来配置。看个例子：

torch.onnx.export(
    model,
    dummy_input,
    "model_dynamic.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "output": {0: "batch_size"}
    }
)

这里我做了三件事：

  • 把输入的 batch 维度（第0维）设为动态，取名 "batch_size"
  • 把输入的高和宽也设为动态
  • 输出的 batch 维度也跟着动态变化

这样导出的 ONNX 就可以接受不同 batch size、不同分辨率的输入了。

注意：不是所有模型都适合全动态。如果你的模型里有全连接层（Linear），它的权重是固定的，输入特征维度变了就会报错。所以一般只把 batch 维度设成动态，高宽维度要看模型结构是否支持。

三、动态轴设置的避坑指南

我在实际项目中踩过不少动态轴的坑，分享几个典型的：

坑1：动态轴名称冲突

如果你有多个输入输出，动态轴的名称不能重复。比如两个输入都把第0维叫 "batch_size"，没问题。但如果一个叫 "batch_size"，另一个也叫 "batch_size"，ONNX 会认为它们是同一个动态维度，这可能导致形状推导错误。

坑2：动态轴与算子兼容性

有些 ONNX 算子不支持动态形状。比如 Resize 算子，在低版本 opset 里要求输出尺寸是常量。我遇到过用 F.interpolate 做上采样，导出时设了动态高宽，结果 Resize 算子报错。解决方案：要么升级 opset，要么改用固定尺寸。

坑3：动态轴导致精度下降

这个比较隐蔽。有些推理引擎在处理动态形状时，会使用更保守的数学优化，导致精度略有下降。我曾在 TensorRT 上遇到过，动态 batch 的推理结果和静态 batch 差了 0.01 的 mAP。如果你的任务对精度极其敏感，建议先跑静态轴验证，再切动态轴。

四、知识结构图：导出流程与决策路径

下面这张图总结了从 PyTorch 模型到 ONNX 导出的完整流程，以及动态/静态轴的选择逻辑：

PyTorch → ONNX 导出流程与动态/静态轴决策 PyTorch 模型 配置核心参数 opset_version / export_params / do_constant_folding 是否需要 动态形状？ 否 静态轴导出 是 动态轴导出 固定形状推理 动态形状推理 建议：生产环境优先用静态轴；需要灵活性的场景（如多batch、多分辨率）才用动态轴 简单高效 灵活但复杂

五、实战建议：什么时候用动态轴？

根据我的经验，以下场景适合用动态轴：

  • 推理服务需要支持可变 batch size：比如线上服务有时来1张图，有时来8张图
  • 输入分辨率不固定：比如目标检测模型要处理不同尺寸的图片
  • NLP 模型序列长度可变：比如文本长度从10到100不等

以下场景建议用静态轴：

  • 生产环境追求极致性能：固定形状可以让推理引擎做更多优化
  • 模型结构对形状敏感：比如有全连接层或固定 reshape 操作
  • 推理引擎对动态形状支持不好：比如某些嵌入式端的引擎

我的个人习惯：先以静态轴导出，跑通整个部署流程。如果后续确实需要动态形状，再改 dynamic_axes。这样排查问题更高效，不会一开始就被动态轴的坑绊住。

好了，关于 torch.onnx.export 的核心参数和动态轴设置，就聊到这里。这些内容看起来琐碎，但每个参数背后都有实际场景的考量。你动手导出几个模型试试，遇到问题再回来对照着看，理解会更深刻。

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