ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch 导出 ONNX 避坑指南:动态轴、算子与数值对拍

PyTorch 导出 ONNX 避坑指南:动态轴、算子与数值对拍

模型在训练侧指标跑得漂亮,交到部署侧却卡在最不起眼的一步——把 PyTorch 导出模型为 ONNX 这件事,看着只有一行torch.onnx.export,真上手才知道坑有多密。我自己遇到过导出后 onnxruntime 输出全零、动态 batch 一换尺寸就报错、某个算子提示不支持、文件导出来好几个 G 加载不动、导出的结果和原模型差了 1e-1 量级。这些没有一个是文档里三言两语能讲清的,基本都是靠一次次对拍、一行行翻报错才定位下来。

这篇按我自己的实操顺序走一遍:导出前该把模型收拾成什么样、参数哪些必须改、动态轴怎么配、报错怎么排、导完怎么验、新版导出路径要不要跟。适合已经能跑通官方 demo、但换成自家模型就各种状况的读者,也适合刚接手模型交付、需要一套可复用检查流程的同学。下面所有代码片段都可以直接改改就用,不依赖特定业务场景。

1. 导出前的三件必做事:先把模型钉死在推理态

很多人拿到模型直接一行 export 甩过去,跑通了就以为完事,结果线上精度不对再回头查,成本翻好几倍。导出前的准备工作决定了后面 80% 的麻烦会不会出现,我现在的习惯是先过三关,一关不过就不往下走。

1.1 eval() 与 no_grad 是两件独立的事,缺一不可

这是个高频误解。model.eval()改的是模块内部的行为开关,让 Dropout 停止随机丢弃、让 BatchNorm 改用运行时统计量而不是当前 batch 的均值方差;torch.no_grad()关的是自动求导图的构建。两者互不替代,只做其中一个都会留下隐患。

只调eval()不关梯度,导出的计算图里可能残留一些训练期才需要的节点,图会变肥,个别算子在转换时会因为梯度相关属性报错。只关梯度不调eval()更危险,因为导出时的前向是在单个 dummy batch上跑的,如果 BatchNorm 还在用当前 batch 统计量,导出的图就会把这个行为固化下来,线上推理时结果会随输入波动而漂移,而且这种偏差在随机输入的对拍里不一定暴露得明显。

我的标准写法是把两件事都写死:

model = MyNet() model.load_state_dict(torch.load("ckpt.pt", map_location="cpu")) model.eval() with torch.no_grad(): torch.onnx.export(model, dummy, "model.onnx", opset_version=17)

提醒一句:eval()之后如果代码里某些层被手动改过training属性,比如某个自定义模块在forward里读取self.training做分支,导出前务必确认走的是推理分支。这类问题排查起来极费时间,因为模型本身不报错,只是结果不对。

1.2 把 forward 里所有训练分支收敛掉

真实项目里的forward往往不干净:有根据self.training走的两套逻辑,有调试用的中间结果打印,有只在训练期生效的辅助 loss 计算,还有一些"暂时留着以后可能用"的实验分支。导出的追踪过程只认实际跑过的那条路径,没走到的分支会被直接丢掉,而这恰恰是问题所在——如果某个分支依赖的条件在导出时恰好成立、线上却相反,导出的图就是错的。

我一般做两件事。第一,在forward里搜一遍self.training、if debug、assert这类关键词,把训练专属逻辑挪出去或者用显式参数控制。第二,给模型准备一份"推理专用"的包装类,只保留真正的推理路径,把辅助输出砍掉,让导出的图尽量干净。

辅助输出这个问题值得多说一句。有些模型forward会返回一个字典或者多个张量,其中一部分推理端根本不用。ONNX 会老老实实把这些输出全部导出,既增大文件体积,也让下游对接时多一层解析。能砍就砍,在包装类里只返回真正需要的那几个。

1.3 dummy input 必须还原线上真实的 dtype 与 shape

dummy input 不是随便torch.randn(1, 3, 224, 224)就完事的。它是导出过程的"探针",追踪器会用它去推每层的形状、类型和设备。有几处必须对齐线上实际:

  • dtype:线上喂的是float32还是半精度,导出就得用什么。用float32导出、线上喂半精度,算子行为可能不一致。
  • shape:至少要和线上常见输入的量级接近。有些模型对小尺寸输入有特殊分支(比如自适应池化的边界处理),用1x1的图去探会探出奇怪的路径。
  • 设备:建议在 CPU 上导出。模型在 GPU 上导出时,某些算子会被绑定设备相关的属性,跨环境加载时容易出问题。真要 GPU 导出,dummy input 和模型必须在同一设备上,否则直接报设备不匹配。

还有一个细节:如果模型有多个输入(比如图像加一个掩码或者一个时间步张量),args要按forward的参数顺序传成一个元组,位置错一个就会导出成完全不同的图,而且不报错。

2. torch.onnx.export 的参数拆解:真正需要动的只有那几个

函数签名参数一大串,第一次看确实唬人。但按我的经验,日常真正需要显式指定的就那么五六个,其余保持默认基本没问题。搞清楚每个参数管什么,比死记一份"万能模板"有用得多。

2.1 args 与 input_names / output_names 的绑定关系

args是位置参数元组,决定追踪器用什么输入去跑前向。input_names和output_names决定导出后图里这些张量叫什么名字。这两组东西必须严格按顺序对应——input_names的第 0 个名字对应args元组的第 0 个元素,第 1 个对应第 1 个,依此类推。

这里有个我踩过的坑:input_names里如果混用了特殊字符或者中文,某些推理端在按名字取输入时会找不到。命名就用最朴素的英文加下划线,input、input_ids、attention_mask这种,别整花活。

output_names同理。如果模型输出是嵌套结构(比如return (logits, (hidden, cell))),导出后输出是扁平化的,名字按遍历顺序依次分配。这时候output_names的数量必须和扁平化之后的张量个数对上,少写一个就会报长度不匹配。我习惯先不写output_names导一次,用onnx.load打印图看看输出到底几个、叫什么,再回头补名字。

import onnx m = onnx.load("model.onnx") print([o.name for o in m.graph.output]) print([i.name for i in m.graph.input])

这两行代码在排查阶段用得非常多,因为图里的实际名字和你在 PyTorch 侧想象的名字经常不一致。

2.2 opset_version:从算子支持矩阵倒着选

opset 版本是导出时最容易拍脑袋决定的参数,也是最容易埋雷的地方。它不是"越高越好",而是"够用且稳定"最好。

原则很简单:先看你的模型用了哪些算子,再查这些算子在哪个 opset 版本才开始被标准支持,取那个版本作为下限。比如注意力相关的结构,在早期 opset 里LayerNormalization并没有成为一等公民,导出时会被拆成一堆基础运算拼起来,图变得臃肿而且数值上容易有细微差别;较新的 opset 里它才作为单个算子导出。再比如GELU,在比较老的版本里同样是被拆开的。

但也不能一味往高版本冲。opset 版本越高,对推理端的要求越高,线上部署环境的推理库版本跟不上就直接加载失败。我一般的做法是:

场景建议 opset理由
传统 CNN 分类/检测模型11 或 13算子简单,低版本兼容性最好
含注意力结构的 Transformer 类模型17相关算子覆盖完整,拆分少
用了较新的插值、采样类算子16 及以上低版本无对应算子
部署环境推理库版本偏旧从 11 起步试兼容优先,跑通再升

一个判断技巧:先按你估计的版本导一次,如果出现Unsupported operator或者日志里出现大量Converting to decomposed之类的警告,就往上升一档再试。如果导出顺利但推理端加载时报版本不支持,就往下降一档。这个二分过程通常两三次就能定下来。

2.3 do_constant_folding、export_params 与 verbose 的实际取舍

export_params=True是默认值,意思把训练好的权重一起写进文件,形成一个自包含的模型。绝大多数场景都要开着,除非你打算把权重和结构分开管理。关掉它导出的只是一个空壳计算图,加载后必须手动灌权重,实际项目里很少这么用。

do_constant_folding=True也是默认值,作用是在导出时把能提前算出来的常量表达式算掉。比如x / 255.0后面跟着一串固定系数的乘加,这些固定部分会被折叠成常量。好处是图更小、推理更快,坏处是折叠之后有些中间节点消失了,出问题时定位会稍微难一点。我日常都开着,只有在排查数值差异怀疑是折叠引起的时候,才临时关掉做对比。

verbose=True是个被低估的参数。打开后能看到导出过程的详细日志,包括每个算子被转成了什么、有没有发生分解、有没有走 fallback 路径。模型导出来结构不对劲的时候,第一件事就是开 verbose 看日志,往往一眼就能看出是哪一层出了问题。

3. dynamic_axes:让同一份文件吃下不同 batch 和变长输入

不配动态轴导出的模型,输入形状是写死的。你用1x3x224x224导出,线上就只能喂1x3x224x224,喂个 batch=4 直接报形状不匹配。这是新手最常撞的墙,也是必须掌握的一块。

3.1 键值语义:名字必须是"导出后"的名字

dynamic_axes的结构是{张量名: {维度序号: 维度别名}}。这里最容易错的是张量名用的是input_names/output_names里定义的名字,不是 PyTorch 里的参数名。你在forward里叫x,但input_names=["input"],那dynamic_axes里就得写"input"。

维度序号从 0 开始数。{0: "batch"}表示第 0 维可变,其余维固定。别名只是个标签,方便阅读,不影响行为。

torch.onnx.export( model, dummy, "model.onnx", input_names=["input"], output_names=["logits"], dynamic_axes={ "input": {0: "batch", 2: "height", 3: "width"}, "logits": {0: "batch"}, }, opset_version=17, )

上面这段的意思是:输入的第 0、2、3 维可变,输出的第 0 维可变。注意输出也要一起配,只配输入不配输出的话,输出形状仍然被钉死在导出时的尺寸上,推理端拿到结果可能被截断或者报错。

还有一处需要留心:输入输出共享的维度别名最好保持一致,都叫batch。某些推理引擎在做形状推导时会利用别名做关联,虽然不写一致也能跑,但一致了排查问题时看着更清楚。

3.2 变长序列与多输入多输出的写法

自然语言类的模型一般至少有输入 ids 和 attention mask 两个输入,序列长度是可变的,那就要把两个输入的序列维都标成动态,同时把 batch 维也标动态。

dynamic_axes = { "input_ids": {0: "batch", 1: "seq"}, "attention_mask": {0: "batch", 1: "seq"}, "logits": {0: "batch", 1: "seq"}, }

图像类的多尺度输入同理,把高宽两个维度标出来就行。但这里有个性能上的现实问题:把空间维度也放开,推理引擎就没法针对固定尺寸做图优化了,速度通常会掉一截。如果你的线上尺寸是固定的(比如永远 640x640),只把 batch 维放开,空间维保持固定,性能会明显好一些。这是我吃过亏之后改过来的习惯——先问清楚线上到底是不是真的需要多尺度。

多输出模型的写法没有额外复杂度,只要保证每个输出名在字典里的键和output_names一一对应就行。漏配某一个输出的动态轴,那个输出就会被固定住,很容易在联调时才发现。

3.3 动态轴打开之后,reshape/view 最先出问题

动态轴一开,模型内部所有依赖具体数值的形状操作都会变成隐患。典型的是x.view(batch_size, -1)这种写法,如果batch_size是从外面传进来的常量,导出时它会被固化成那个数,动态输入进来就崩。

正确做法是用-1让框架自己推,或者用x.flatten(1)这种跟维度无关的表达。类似的还有:

  • x.view(x.size(0), 256, -1)比x.view(4, 256, -1)安全得多,前者依赖的是运行时尺寸。
  • 用torch.Tensor.expand时如果参数是硬编码的 batch,同样会被固化。
  • 循环里按range(x.shape[0])展开的逻辑,导出后会被展开成固定次数的重复结构,序列长度一变就错。

我排查这类问题的办法是:导出时故意用两个不同 batch 的 dummy input 各导一次,对比生成的两个 onnx 文件的图结构,如果某处节点的常量参数不一样,那里就有硬编码。这个方法笨但很有效。

补一句实操经验:新版 PyTorch 对动态形状的推导比早几年强很多,但仍然建议导出后用至少两个不同的 batch 实测,一次 batch=1、一次 batch=4,两个都跑通才算过关。只测一个尺寸等于没测。

4. 导出失败与数值对不齐的排查链路

导出报错分两类,一类是直接抛异常导不出来,另一类是导出来了但结果对不上。后者更难查,因为没有任何报错信息,只能靠对拍。

4.1 不支持算子的定位与 symbolic 注册

遇到Unsupported operator的时候,日志通常会带上算子名和出现的位置。第一步是确认这个算子是不是真的没有对应实现——去查一下当前 opset 下有没有等价算子,很多时候升个 opset 就解决了。

如果确认是自定义算子或者确实没实现,有两条路。一条是把自定义算子在前向里用基本算子重写一遍,牺牲一点性能换可移植性,这条路最稳。另一条是注册自定义 symbolic,把这层算子映射到你推理端支持的自定义域算子。

from torch.onnx import register_custom_op_symbolic import torch def my_op_symbolic(g, self, other): return g.op("custom_domain::MyOp", self, other) register_custom_op_symbolic("myns::my_op", my_op_symbolic, 17)

注册完之后推理端必须也支持这个自定义域算子,否则加载时会报未知算子。跨团队协作时这一点要提前对齐,光导出端搞定没用。

还有一种隐蔽情况:算子本身支持,但某个参数组合不支持。日志里会写类似"attribute not supported"的提示,这种通常需要改写调用方式绕开。

4.2 双端对拍:checker + onnxruntime 的最小验证脚本

导出成功不等于导出正确。我现在的固定动作是跑一个最小对拍脚本,三件事一起做:校验模型合法性、跑推理、和 PyTorch 结果比对。

import numpy as np import onnx import onnxruntime as ort import torch # 1. 结构校验 onnx_model = onnx.load("model.onnx") onnx.checker.check_model(onnx_model) # 2. 构建推理会话 sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"]) # 3. 同输入双端对比 x = torch.randn(2, 3, 224, 224) with torch.no_grad(): ref = model(x).numpy() out = sess.run(None, {"input": x.numpy()})[0] diff = np.abs(ref - out) print("max abs diff:", diff.max()) np.testing.assert_allclose(ref, out, rtol=1e-3, atol=1e-5)

onnx.checker.check_model能抓出图结构本身的问题,比如节点缺输入、类型不匹配、opset 声明和实际使用不一致。这个检查很便宜,几毫秒的事,但能挡掉一批低级错误,建议每次导出都跑。

容差怎么定?rtol=1e-3, atol=1e-5是我常用的起点。分类模型看的是 argmax 是否一致,1e-4量级的差异完全无影响;回归类、检测框坐标类的模型就得盯紧绝对值,1e-2的偏差都可能让框偏移肉眼可见。如果差异超过预期,先别急着怀疑导出,把 PyTorch 侧同一输入跑两遍(一次带梯度一次不带)看看是不是本身就有一致性问题,能排除掉一批误判。

对拍的时候输入别只用一个。我一般跑三组:全零、全一、随机。全零输入能暴露某些除零或者归一化相关的边界问题,随机输入覆盖常规情况。

4.3 一张常见报错对照表

下面这些是我这些年反复遇到的,按出现频率排:

报错或现象大概率原因处理方向
输出全零或全为常数忘了eval(),BatchNorm/Dropout 行为异常补model.eval()并用no_grad重导
换 batch 尺寸报 shape 不匹配未配dynamic_axes或只配了输入没配输出输入输出动态轴一起配
Unsupported operatoropset 太低或用了自定义算子升 opset;仍不行则重写或注册 symbolic
加载时提示模型过大超过单文件容量上限改用外部权重文件存储
输出形状对但数值差很多动态轴下存在硬编码 reshape排查view/expand的常量参数
推理端按名字取输入失败input_names含特殊字符改成纯英文下划线命名
导出过程极慢或内存爆掉dummy input 过大或模型在训练模式减小 dummy 尺寸,确认推理态

这张表不能覆盖所有情况,但能挡掉八成常见问题。剩下两成基本都得靠开 verbose 看日志、逐层对比中间输出来定位。

5. 导出之后的收尾活儿:简化、拆权重与端上实测

导出一个能跑的文件只是中间状态,离能交付还有几步。很多人导完就直接扔给下游,结果下游反馈文件太大、加载慢、启动要十几秒,回头再优化又是一轮。

5.1 onnx-simplifier 什么时候值得跑

导出的图里经常带一堆冗余节点:常量折叠没吃干净的、shape 推导留下的辅助算子、恒等映射之类的。onxx-simplifier 这类工具会做常量传播和算子融合,把图压小。

什么时候值得跑?我的判断标准是:如果 verbose 日志里出现了大量分解产生的琐碎节点,或者图节点数比模型层数高出一个数量级,就跑一遍。跑完之后节点数通常能降下来不少,加载速度和首帧延迟都会有改善。

但也不是没风险。简化工具在做 shape 推断时会用固定尺寸去跑一遍,如果模型带动态轴,简化后有可能把动态性给"优化"掉。跑完必须重新做一遍 4.2 的对拍,而且要用不同于简化时用的输入尺寸去验证。我自己就遇到过一次,简化之后 batch 维被固化成 1,单测通过、线上多 batch 直接崩。

5.2 超过 2GB 的模型怎么处理

单文件格式有个硬上限,超过之后就存不下。大模型必须改成"结构文件 + 外部权重文件"的形式。

import onnx onnx.save_model( onnx_model, "model.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="model.onnx.data", size_threshold=1024, )

保存之后是两个文件,加载时必须放在同一目录下,路径关系写错就报找不到权重。新版torch.onnx.export也有对应参数可以直接导出外部权重,省掉二次保存,但参数名在不同版本间有变动,用之前先查一下当前版本的签名。

交付这类模型时,我习惯把权重文件和结构文件打包在一起,附带一个加载示例。光发一个.onnx给下游,对方拿到手发现加载失败,来回沟通的成本比多写两行说明高得多。

5.3 端上实测要看哪几个指标

对拍通过只说明数值对,性能还得单独看。我一般关注四个数:首帧延迟、稳态单帧延迟、峰值内存、模型加载耗时。

首帧延迟高往往是算子初始化或者图优化在第一次推理时做的,属于一次性成本;稳态延迟才是线上真正关心的。峰值内存决定能不能塞进目标设备。加载耗时在服务化场景下很关键,冷启动慢会直接影响扩容速度。

测的时候用真实的推理库、真实的输入尺寸,别用随机尺寸凑数。同一个模型在固定尺寸和动态尺寸下的性能差异可能到百分之二三十,用错配置测出来的数没有参考价值。

6. PyTorch 2.x 的新导出路径:dynamo 值不值得上

近两年导出机制在变,新版本里torch.onnx.export背后的实现已经换了一套,内部走的是基于图捕获的新路径,另外还单独提供了新的导出接口。这块如果不了解,遇到问题时会很懵——明明代码没改,换个版本行为就不一样了。

6.1 新路径和老路径的差别在哪

老路径是追踪式的:拿 dummy input 跑一遍前向,记录实际执行过的算子。优点是简单直接,缺点是对控制流无能为力——if/for依赖输入值的分支,只会记录走的那一条,另一条直接丢失。

新路径是图捕获式的:先把前向的逻辑捕获成一张计算图,再统一转成目标格式。对控制流的处理好得多,能表达条件分支和循环,对动态形状的支持也更自然。代价是捕获过程对代码有一定要求,遇到不支持的写法会直接报错,而不是像老路径那样"悄悄记错"。

新版export默认走新路径,遇到问题可以显式传dynamo=False回退到老路径:

torch.onnx.export(model, dummy, "model.onnx", dynamo=False, opset_version=17)

动态形状的写法在新路径里也变了,不再用dynamic_axes,而是通过dynamic_shapes传入形状约束。两套 API 不能混用,混了会报参数冲突。

6.2 老项目要不要迁:一个判断标准

我的判断很简单:现有导出流程能稳定产出、验证通过、线上没出过问题的,不迁。迁移的收益主要在控制流复杂、动态形状多的模型上,如果你的模型是标准的前馈结构,老路径完全够用,而且周边工具链、排查经验都更成熟。

需要迁的信号有几个:模型里有大量数据依赖的分支(比如检测后处理里的 NMS 逻辑、变长序列的循环);动态尺寸场景复杂,dynamic_axes表达不了;现有路径下反复出现"追踪到了错误分支"的问题。这几种情况下新路径的优势明显。

真要迁,别一次性全量替换。我的做法是两套并行跑一段时间,同一批测试输入分别走两条路导出,对比输出和性能,稳定之后再切换。中间留一个开关,出问题能立刻回退,这个保险值得留。

最后分享一个小技巧:不管走哪条路径,我都会在导出脚本末尾加一段自动校验——加载刚导出的文件、跑一组固定输入、和当前 PyTorch 输出比对,不一致就让脚本非零退出。这样 CI 里能直接拦住问题,不用等到联调阶段才发现导出的模型是坏的。这套机制上线之后,模型交付阶段的返工次数少了很多。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进