模型在训练侧指标跑得漂亮交到部署侧却卡在最不起眼的一步——把 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_locationcpu)) model.eval() with torch.no_grad(): torch.onnx.export(model, dummy, model.onnx, opset_version17)提醒一句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 与 shapedummy 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_paramsTrue是默认值意思把训练好的权重一起写进文件形成一个自包含的模型。绝大多数场景都要开着除非你打算把权重和结构分开管理。关掉它导出的只是一个空壳计算图加载后必须手动灌权重实际项目里很少这么用。do_constant_foldingTrue也是默认值作用是在导出时把能提前算出来的常量表达式算掉。比如x / 255.0后面跟着一串固定系数的乘加这些固定部分会被折叠成常量。好处是图更小、推理更快坏处是折叠之后有些中间节点消失了出问题时定位会稍微难一点。我日常都开着只有在排查数值差异怀疑是折叠引起的时候才临时关掉做对比。verboseTrue是个被低估的参数。打开后能看到导出过程的详细日志包括每个算子被转成了什么、有没有发生分解、有没有走 fallback 路径。模型导出来结构不对劲的时候第一件事就是开 verbose 看日志往往一眼就能看出是哪一层出了问题。3. dynamic_axes让同一份文件吃下不同 batch 和变长输入不配动态轴导出的模型输入形状是写死的。你用1x3x224x224导出线上就只能喂1x3x224x224喂个 batch4 直接报形状不匹配。这是新手最常撞的墙也是必须掌握的一块。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_version17, )上面这段的意思是输入的第 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 实测一次 batch1、一次 batch4两个都跑通才算过关。只测一个尺寸等于没测。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, rtol1e-3, atol1e-5)onnx.checker.check_model能抓出图结构本身的问题比如节点缺输入、类型不匹配、opset 声明和实际使用不一致。这个检查很便宜几毫秒的事但能挡掉一批低级错误建议每次导出都跑。容差怎么定rtol1e-3, atol1e-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_dataTrue, all_tensors_to_one_fileTrue, locationmodel.onnx.data, size_threshold1024, )保存之后是两个文件加载时必须放在同一目录下路径关系写错就报找不到权重。新版torch.onnx.export也有对应参数可以直接导出外部权重省掉二次保存但参数名在不同版本间有变动用之前先查一下当前版本的签名。交付这类模型时我习惯把权重文件和结构文件打包在一起附带一个加载示例。光发一个.onnx给下游对方拿到手发现加载失败来回沟通的成本比多写两行说明高得多。5.3 端上实测要看哪几个指标对拍通过只说明数值对性能还得单独看。我一般关注四个数首帧延迟、稳态单帧延迟、峰值内存、模型加载耗时。首帧延迟高往往是算子初始化或者图优化在第一次推理时做的属于一次性成本稳态延迟才是线上真正关心的。峰值内存决定能不能塞进目标设备。加载耗时在服务化场景下很关键冷启动慢会直接影响扩容速度。测的时候用真实的推理库、真实的输入尺寸别用随机尺寸凑数。同一个模型在固定尺寸和动态尺寸下的性能差异可能到百分之二三十用错配置测出来的数没有参考价值。6. PyTorch 2.x 的新导出路径dynamo 值不值得上近两年导出机制在变新版本里torch.onnx.export背后的实现已经换了一套内部走的是基于图捕获的新路径另外还单独提供了新的导出接口。这块如果不了解遇到问题时会很懵——明明代码没改换个版本行为就不一样了。6.1 新路径和老路径的差别在哪老路径是追踪式的拿 dummy input 跑一遍前向记录实际执行过的算子。优点是简单直接缺点是对控制流无能为力——if/for依赖输入值的分支只会记录走的那一条另一条直接丢失。新路径是图捕获式的先把前向的逻辑捕获成一张计算图再统一转成目标格式。对控制流的处理好得多能表达条件分支和循环对动态形状的支持也更自然。代价是捕获过程对代码有一定要求遇到不支持的写法会直接报错而不是像老路径那样悄悄记错。新版export默认走新路径遇到问题可以显式传dynamoFalse回退到老路径torch.onnx.export(model, dummy, model.onnx, dynamoFalse, opset_version17)动态形状的写法在新路径里也变了不再用dynamic_axes而是通过dynamic_shapes传入形状约束。两套 API 不能混用混了会报参数冲突。6.2 老项目要不要迁一个判断标准我的判断很简单现有导出流程能稳定产出、验证通过、线上没出过问题的不迁。迁移的收益主要在控制流复杂、动态形状多的模型上如果你的模型是标准的前馈结构老路径完全够用而且周边工具链、排查经验都更成熟。需要迁的信号有几个模型里有大量数据依赖的分支比如检测后处理里的 NMS 逻辑、变长序列的循环动态尺寸场景复杂dynamic_axes表达不了现有路径下反复出现追踪到了错误分支的问题。这几种情况下新路径的优势明显。真要迁别一次性全量替换。我的做法是两套并行跑一段时间同一批测试输入分别走两条路导出对比输出和性能稳定之后再切换。中间留一个开关出问题能立刻回退这个保险值得留。最后分享一个小技巧不管走哪条路径我都会在导出脚本末尾加一段自动校验——加载刚导出的文件、跑一组固定输入、和当前 PyTorch 输出比对不一致就让脚本非零退出。这样 CI 里能直接拦住问题不用等到联调阶段才发现导出的模型是坏的。这套机制上线之后模型交付阶段的返工次数少了很多。