3、PyTorch转ONNX：torch.onnx.export核心参数详解与常见报错

说实话，PyTorch转ONNX这一步，是模型部署里最容易出幺蛾子的环节。我自己刚入行那会儿，光一个torch.onnx.export就卡了整整两天。后来才发现，其实就是几个参数没搞明白。

今天咱们就把这个函数掰开揉碎了讲。你只要掌握了这几个核心参数，90%的转换问题都能自己解决。

3.1 核心参数一览

先看一个最常用的调用方式：

import torch

# 假设你有一个训练好的模型
model = MyModel()
model.eval()

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

# 导出 ONNX
torch.onnx.export(
    model,                     # 要导出的模型
    dummy_input,               # 示例输入
    "model.onnx",              # 输出路径
    export_params=True,        # 是否导出参数
    opset_version=11,          # ONNX算子集版本
    do_constant_folding=True,  # 是否做常量折叠
    input_names=["input"],     # 输入节点名称
    output_names=["output"],   # 输出节点名称
    dynamic_axes={             # 动态轴定义
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)

嗯，这里每个参数都有讲究。我一个个说。

3.2 export_params：要不要带权重？

这个参数默认是True。说白了就是：导出的ONNX文件里，要不要把模型的权重参数也打包进去。

我个人习惯是：

  • 部署用：必须True。不然你拿到的就是个空壳子，没法推理。
  • 做模型结构分析：可以False。文件小，加载快，只看网络结构。
小技巧：如果你发现导出的ONNX文件特别小（比如只有几KB），先检查一下是不是忘了设export_params=True。我在项目里就犯过这个低级错误，排查了半天。

3.3 opset_version：算子版本选哪个？

这个参数特别容易让人纠结。ONNX每个版本都会新增或废弃一些算子。

我的建议是：

  • 通用场景：选11。兼容性好，大部分推理引擎都支持。
  • 需要新特性：选13或15。比如你想用If算子或者Loop算子。
  • 别选太新：比如18、19。很多推理框架还没跟上，导出容易报错。

为什么会这样？因为ONNX版本和推理引擎的版本是绑定的。你选了个太新的opset，推理引擎不认识，那就只能报错了。

注意：我曾经在项目中用了opset=17，结果部署到TensorRT上直接报"Unsupported operator"。后来降到11，一切正常。所以，除非你有特殊需求，否则老老实实用11或13。

3.4 do_constant_folding：要不要做常量折叠？

这个参数默认是True。它的作用是：把模型里那些固定不变的运算，提前算好，直接替换成常量。

举个例子：

# 如果你的模型里有这样的操作
x = torch.ones(3, 3) * 2  # 这个结果永远是 [[2,2,2],[2,2,2],[2,2,2]]

开启常量折叠后，ONNX导出时就会直接把这个结果写死，而不是保留一个乘法算子。

好处很明显：模型变小了，推理速度也快了。但有个坑——如果你的模型里有动态形状相关的操作，常量折叠可能会把形状信息搞丢。

经验之谈：我一般先开True导出。如果报错说形状不匹配，再改成False试试。大多数情况下，True是没问题的。

3.5 input_names / output_names：命名很重要

这两个参数看似简单，但命名不规范会带来大麻烦。

我见过有人这么写：

input_names=["0"],  # 别这么干！
output_names=["1"]

结果部署的时候，推理引擎报错说找不到节点。为什么？因为很多推理框架内部会用数字做索引，你再用数字命名，就冲突了。

我的命名规范：

  • 用有意义的英文单词：input、output、image、features
  • 不要用数字开头
  • 不要包含特殊字符（下划线可以）

3.6 dynamic_axes：动态形状怎么配？

这是最容易出错的地方。先看一个典型配置：

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

这里的意思是：

  • 输入的batch_size是可变的
  • 输入的高度和宽度也是可变的
  • 输出的batch_size跟着输入变

但有个坑：如果你把所有维度都设成动态，ONNX导出时可能会报错。因为有些算子（比如Reshape）要求某些维度是固定的。

避坑指南：我曾经把一个3D卷积模型的全部4个维度都设成动态，结果导出报错"Unsupported dynamic dimension"。后来只把batch_size设成动态，其他维度固定，就顺利通过了。

我的建议是：

  • 只把确实需要变化的维度设成动态
  • batch_size几乎总是动态的
  • 图像尺寸如果固定，就别设动态
  • 如果必须动态，先测试几个不同的尺寸，确保都能跑通

3.7 常见报错与解决方案

这里列几个我遇到最多的报错：

报错信息 原因 解决方案
RuntimeError: ONNX export failed 模型里有不支持的算子 检查模型，替换成ONNX支持的算子
RuntimeError: Unsupported operator opset版本太低，不支持某些算子 升级opset版本，或者替换算子
RuntimeError: Input shape mismatch dynamic_axes配置错误 检查动态轴定义，确保和模型匹配
RuntimeError: Constant folding failed 常量折叠时遇到动态形状 关闭常量折叠：do_constant_folding=False
调试技巧：遇到报错时，先看报错信息里的行号。ONNX导出会告诉你是在哪个算子出错的。然后去模型里找到对应的代码，看看是不是用了不支持的函数。

3.8 一个完整的导出示例

最后，给你一个我常用的模板：

import torch
import torch.onnx

def export_to_onnx(model, dummy_input, output_path):
    model.eval()
    
    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        export_params=True,
        opset_version=11,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {0: "batch_size"},
            "output": {0: "batch_size"}
        }
    )
    
    print(f"模型已导出到: {output_path}")

# 使用示例
model = YourModel()
dummy = torch.randn(1, 3, 224, 224)
export_to_onnx(model, dummy, "your_model.onnx")

这个模板我用了两年多，基本没出过问题。你直接拿去用就行。

嗯，关于torch.onnx.export的核心参数，今天就讲到这里。下一节咱们聊聊ONNX模型的可视化和调试工具，到时候你会发现，很多问题其实看一眼图就明白了。

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