3、模型导出实战（PyTorch篇）：torch.onnx.export核心参数详解、动态轴与静态轴设置、导出常见错误与修复

好，咱们进入实战环节。

前面聊了那么多ONNX的理论，说实话，不亲手导一次模型，你永远不知道坑在哪。我个人习惯是，先把PyTorch模型跑通，再考虑导出。今天这一章，咱们就死磕 torch.onnx.export 这个函数。

你想想看，一个模型从PyTorch到ONNX，本质上就是一次「翻译」。翻译得好不好，全看参数调得对不对。我刚开始接触ONNX时，也踩过不少雷，导出的模型要么精度对不上，要么推理直接崩。嗯，今天就把这些经验全抖出来。

3.1 torch.onnx.export 核心参数详解

这个函数，说白了就干一件事：把PyTorch模型的计算图，转换成ONNX的图结构。它的签名很长，但常用的核心参数其实就几个。

torch.onnx.export(
    model,               # 要导出的PyTorch模型
    args,                # 模型的输入张量（或一组输入）
    f,                   # 输出文件路径，比如 "model.onnx"
    export_params=True,  # 是否导出模型权重
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True, # 是否做常量折叠优化
    input_names=['input'],     # 输入节点名称
    output_names=['output'],   # 输出节点名称
    dynamic_axes=None          # 动态轴设置，后面重点讲
)

咱们一个一个拆开说。

3.1.1 model 和 args

model 就是你的PyTorch模型实例，记得要设成 model.eval() 模式。为什么？因为训练模式下的Dropout、BatchNorm行为跟推理不一样，导出的ONNX会乱掉。我在项目中就见过有人忘了这步，结果导出的模型推理结果全是NaN，排查了半天。

args 是模型的输入。可以是单个张量，也可以是一个元组。这里有个小技巧：你给的输入张量，ONNX会用它来「跑一次」模型，从而追踪计算图。所以这个输入的形状和数据类型，最好跟你实际推理时保持一致。

3.1.2 opset_version

这个参数我建议你设成11或更高。为什么？因为低版本的ONNX算子集，很多PyTorch的高级操作不支持。比如 torch.nn.functional.grid_sample，在opset 9里就没有对应的算子，导出时会报错。我个人习惯用opset 11，兼容性和功能都比较平衡。

注意： opset版本不是越高越好。版本太高，一些老的推理引擎（比如某些嵌入式设备上的ONNX Runtime）可能不支持。导出前先确认目标平台的ONNX Runtime版本。

3.1.3 do_constant_folding

这个参数默认是True。它会把计算图中那些「输入全是常量」的子图，提前计算好，然后替换成一个常量节点。说白了就是做静态优化，减少推理时的计算量。

但要注意：如果你的模型里有动态形状的操作，常量折叠可能会把一些本不该折叠的节点也给折叠了，导致导出的模型形状信息丢失。我遇到过这种情况，导出的ONNX模型在推理时形状死活不对，最后发现是常量折叠搞的鬼。遇到这种问题，可以试着把 do_constant_folding=False 试试。

3.2 动态轴与静态轴设置

这是整个导出过程中最容易出问题的地方，也是面试常问的点。

3.2.1 什么是动态轴？

静态轴，就是模型输入/输出的某个维度，在推理时是固定不变的。比如你训练时batch size固定为32，那第0维就是静态的。

动态轴，就是允许某个维度在推理时变化。比如你想让模型支持任意batch size的输入，那第0维就得设成动态的。

为什么要区分？因为ONNX的图结构是静态的，如果你不告诉它某个维度可以变化，它就会把这个维度当成固定值写死。到时候你传一个不同batch size的输入，推理引擎直接报错。

3.2.2 如何设置动态轴？

通过 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=11
)

上面这个例子，我把输入的batch size、height、width都设成了动态的。注意，height和width也设成动态，是因为我想让模型支持任意尺寸的图片输入。

小技巧： 如果你不确定哪些维度需要动态，可以先全部设成静态导出，然后在推理时用不同形状的输入测试。如果报错，再回来改。我刚开始就是这么干的，省得一次设太多反而出问题。

3.2.3 动态轴的坑

嗯，这里要注意。动态轴虽然灵活，但会带来两个问题：

  • 性能下降： 动态形状会让推理引擎无法做某些内存优化，推理速度可能会慢10%-20%。
  • 算子兼容性： 有些ONNX算子不支持动态形状，比如Reshape如果目标形状是动态的，可能会导出失败。

我曾经在一个OCR模型里，把文本序列长度设成了动态轴。结果导出时，torch.nn.CTCLoss 相关的算子死活不支持动态形状，最后只能把序列长度固定成最大长度，用padding的方式解决。虽然浪费了一点计算资源，但至少能跑起来。

3.3 导出常见错误与修复

这部分是我最想分享的。说实话，导出ONNX时遇到的错误，90%都可以归为以下几类。

3.3.1 错误一：Unsupported ONNX ops

现象： 导出时报错，说某个PyTorch操作没有对应的ONNX算子。

原因： 你用的PyTorch操作太新，或者太冷门，ONNX的算子集里没有。

修复：

  • 升级opset版本（比如从9升到11）。
  • 用 torch.onnx.register_custom_op_symbolic 自定义映射。
  • 或者，干脆用PyTorch里更基础的操作组合来实现。比如 torch.nn.functional.leaky_relu 如果导出失败，可以手动写成 torch.where(x > 0, x, negative_slope * x)。

3.3.2 错误二：Shape mismatch

现象： 导出成功，但推理时输入形状不对，报shape mismatch。

原因： 静态轴和动态轴设置冲突。或者模型内部有 torch.reshape 操作，目标形状依赖于动态的输入。

修复：

  • 检查 dynamic_axes 设置，确保所有需要变化的维度都标记了。
  • 如果模型内部有 torch.reshape，尽量用 torch.view 或 torch.flatten 替代，这些操作对动态形状更友好。

3.3.3 错误三：Constant folding 导致的精度问题

现象： 导出的ONNX模型，推理结果跟PyTorch原模型对不上。

原因： 常量折叠把一些本不该折叠的节点给优化掉了，导致计算图变了。

修复： 设置 do_constant_folding=False，然后重新导出。如果问题解决了，那就是常量折叠的锅。这时候你可以手动分析哪些子图可以安全折叠，再选择性开启。

3.3.4 错误四：模型包含控制流

现象： 导出时卡住，或者导出的ONNX模型推理结果完全不对。

原因： PyTorch模型里有 if 语句或 for 循环，这些控制流在静态图中很难表示。

修复：

  • 尽量用 torch.where 替代 if 语句。
  • 如果必须用循环，考虑用 torch.jit.script 先把模型转成TorchScript，再从TorchScript导出ONNX。这条路我走过，虽然麻烦点，但能解决大部分控制流问题。

核心总结： 导出ONNX，本质上就是一次「图追踪」。你给一个输入，它跑一遍，把路径上的所有操作记录下来。所以，确保你的模型在导出时是确定性的，没有随机性，没有动态控制流，没有未初始化的参数。做到这三点，导出基本不会出大问题。

好了，这一章的内容就到这里。下一章咱们会聊TensorFlow模型的导出，到时候你会发现，不同框架的导出思路其实大同小异，但细节上各有各的坑。嗯，到时候再细说。

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