★ FEATURED ARTICLE 模型训练完只是半成品能不能落到业务线上八成卡在导出这一步。这两年我经手的推理部署项目里十有八九要把 PyTorch 训练出来的权重转成 ONNX再交给后端的推理引擎去跑。PyTorch导出模型为onnx这件事表面上看就是一行torch.onnx.export的调用真上手才知道坑密得像筛子动态 batch 导出来变成写死 1、自定义算子直接报不支持、数值对不齐差出小数点后三位、导出的模型在本地跑得好好的到了线上就崩。这篇东西写给已经会训模型、正准备把模型往外交付的人也写给被导出问题折腾过几轮、想把流程做扎实的人。我会把方案选型、关键参数、验证手段、排错路径和工程化落地一次讲透尽量让你少走我走过的弯路。1. 导出这件事到底在解决什么问题1.1 ONNX 在推理链路里的位置先把概念摆正。ONNX 不是模型格式的终点它是一个中间表示层。你在 PyTorch 里定义的计算图本质上是 Python 对象加上一堆 eager 执行的算子调用而 ONNX 描述的是一张静态的、与框架无关的计算图节点是标准算子边是张量附带形状和类型信息。训练框架负责产出这张图推理引擎负责消费这张图两者解耦这就是整个生态存在的理由。理解这一点很关键因为它决定了你导出的目标不是存一份权重而是翻译一张图。权重会被塞进图里变成 initializer结构会被翻译成节点而翻译过程中任何 PyTorch 特有、ONNX 标准里没有对应物的写法都会变成麻烦。我见过太多人把导出当成torch.save的替代品结果导出来的模型根本跑不动或者跑出来的结果和原模型对不上。从工程角度看ONNX 带来的收益集中在三块。一是跨框架、跨语言部署C、C#、Java、Python 都能加载同一份文件二是推理引擎可以做图层面的优化算子融合、常量折叠、内存复用这些手段在 eager 模式下是没法做的三是交付边界清晰你把一个自包含的文件交给下游对方不需要装整套训练框架。这三点决定了它在生产环境里的位置也决定了你对导出质量的容忍度应该有多低。1.2 什么样的项目真的需要导出不是所有项目都值得走这条路。如果你的模型只在训练脚本里做评估或者服务端本来就是 Python 环境、QPS 也不高直接加载.pt用 PyTorch 跑反而省事。导出带来的额外成本是实实在在的你要维护一套对齐验证要处理算子兼容要跟着框架版本升级重新适配。真正需要导出的场景我归纳成几类。第一类是端侧部署手机、嵌入式板子、浏览器这些环境压根装不了完整训练框架。第二类是高吞吐服务推理引擎在同样的硬件上能压出更高的吞吐和更低的尾延迟尤其是 batch 推理场景。第三类是异构硬件专用加速卡通常只认标准图格式你得先把图交出去。第四类是多团队协作算法团队交模型、工程团队做服务中间需要一个格式契约ONNX 就是那个契约。反过来说如果你的模型里大量依赖 Python 控制流、动态数据结构、运行时才知道的 shape那导出的代价会非常高可能改模型结构比导出本身还费劲。这种时候要先评估是改造模型划算还是保留原框架直接部署划算。我个人的判断线是如果为了让模型可导出需要改动超过三成的核心逻辑就先别急着导先看看业务侧的真实瓶颈在哪里。1.3 导出前先想清楚的三件事动手前有三个问题必须先有答案否则后面全是返工。第一个问题是目标运行时的算子覆盖范围。不同的推理引擎支持的 opset 版本和算子集不完全一样有些引擎对某些算子的高版本实现支持滞后。你要先确认下游能吃哪个 opset再倒推导出时用什么版本。这个顺序不能反先导出再问对方支持什么大概率要重来。第二个问题是输入的 shape 语义。你的模型是要固定 batch 还是动态 batch图像尺寸是固定的还是可变的序列长度是否有上限。这些决定了dynamic_axes怎么配也决定了后续能不能做 shape 推断优化。很多线上事故的根源就是导出时把 batch 写死成 1压测的时候才发现根本没法批量推理。第三个问题是数值精度要求。ONNX 默认用 float32如果你的模型里有混合精度训练出来的权重或者对数值敏感的操作导出前后的误差必须量化评估。我一般会在导出前先定一个误差阈值比如分类任务 top-1 一致率必须 100%回归任务最大绝对误差不超过 1e-4达不到就打回重做。2. 方案选型三条导出路径怎么挑2.1 传统 torch.onnx.export 路径这是最主流、文档最全的一条路。它基于 TorchScript 的 tracing 机制把模型跑一遍记录下实际执行的算子序列再翻译成 ONNX 图。优点是成熟稳定、兼容性好绝大多数量化、视觉、NLP 模型都能走通缺点是 tracing 只记录跑过的那条路径控制流里没走到的分支会直接丢失。用它的时候有个前提容易被忽略模型必须处于 eval 模式。BatchNorm 在 train 模式下会用当前 batch 的统计量Dropout 会随机置零trace 出来的图和推理时的行为完全不是一回事。我见过导出后精度暴跌的案例查了半天发现是忘了model.eval()。这一行代码不是可选项是硬性要求。另一个隐性坑是 tracing 的输入。你传进去的 dummy input 不只是用来触发前向它的 shape 和 dtype 会直接写进图里。如果你传的是torch.randn(1, 3, 224, 224)那导出的图默认就是 batch1、分辨率 224。想要动态必须在dynamic_axes里显式声明不能指望它自动推断。2.2 TorchDynamo 后端路径新版本 PyTorch 提供了基于 TorchDynamo 的导出方式走的是字节码分析而不是 tracing对控制流的支持明显更好。Python 层的 if/for 能在一定程度上被正确捕获而不是简单地按 trace 到的路径固化。这对含有条件分支的模型是实打实的改善。但这条路目前的状态是能用但要挑场景。它对模型写法的要求更严格一些动态的 Python 操作、自定义容器、运行时属性访问仍然会触发图断裂断裂的部分会回退到 Python 执行导出结果可能包含不完整的图或者直接报错。我的建议是新项目、模型结构干净、框架版本较新可以优先试 Dynamo老项目、结构复杂、依赖大量自定义逻辑先用传统路径保底。实际选型时我会两个都试一遍对比导出成功率和验证误差哪个稳用哪个。不要迷信新路径导出这件事的评判标准只有一个下游能不能稳定跑出正确结果。2.3 脚本化与 tracing 的取舍除了上面两条torch.jit.script也是一条路先把模型编译成 TorchScript 再导出。它对控制流的支持比 trace 好但要求模型代码符合 TorchScript 的语法子集很多 Python 惯用写法会直接编译失败。改造工作量往往不小。三条路径对比下来我的经验是这样纯前馈结构、无复杂控制流传统 export 最快最稳有 if/while 但结构规整Dynamo 值得一试模型里全是 Python 动态逻辑先考虑重构别硬导。下面这张表是我在实际项目里总结的选型参考。路径控制流支持改造工作量稳定性适用场景torch.onnx.exporttrace弱只记录实际路径低高常规 CNN、Transformer 推理模型Dynamo 导出较强中中新版本框架、结构规整的新项目jit.script 导出较强高中控制流多、可接受改代码手工重写图完全可控极高高极小模型、特殊算子2.4 版本矩阵与依赖锁定导出失败的原因里版本不匹配能占三成。PyTorch 版本、ONNX 版本、opset 版本、推理引擎版本四者之间有一张隐性的兼容矩阵。你本地用新框架导出的图下游用老引擎加载报个unsupported op再正常不过。我的做法是把这套版本写进项目的依赖清单并且和下游对齐后锁死。具体要锁的包括训练侧框架版本、导出时用的 opset、onnx 检查工具版本、推理引擎版本。这四样任何一样变动都要重跑一遍对齐验证。别嫌麻烦线上因为版本漂移导致模型加载失败的排查成本比这套流程高一个数量级。还有个小细节导出脚本里最好显式打印各组件版本写进产物旁边的元信息文件。半年后你回头查一个老模型的导出参数这份记录能救命。3. 动手之前必须搞懂的核心概念3.1 计算图、动态轴与 shape 语义ONNX 图里的每个张量都有 shape但 shape 有两种静态的和符号化的。静态 shape 就是一个确定的数字符号化 shape 是一个带名字的维度运行时可以是任意值。dynamic_axes干的事就是把指定维度从静态改成符号化。这里有个认知误区要纠正dynamic_axes不是允许动态而是声明这一维是动态的。你不声明它就是一个固定值你声明了下游才能按变长输入来做内存规划和优化。两者的差别在部署时会被放大尤其是 batch 维度和序列长度维度。配置的时候要按维度索引来写容易写错。对于一个四维输入[N, C, H, W]索引 0 是 batch1 是通道2 和 3 是空间维度。如果你写成了{1: batch}那声明的就是通道维动态运行时喂不同通道数的输入模型行为会完全错乱。这种错误不会在导出时报错只会在推理时给你一个莫名其妙的 shape mismatch。我的习惯是给每个动态维度起有意义的名字batch 就叫batch高度宽度就叫height、width序列长度就叫seq_len。名字本身不参与计算但它是给人和给工具看的契约出了问题一眼能定位。3.2 opset 版本与算子支持opset 是 ONNX 的算子集版本每个版本会新增算子、修改已有算子的定义。选版本的核心原则是够用且下游支持不是越高越好。高版本可能引入了下游引擎还没实现的新算子低版本又可能缺你需要的算子。我一般的选法是先查下游推理引擎的文档确认它支持的最高 opset再查你模型里用到的算子在这个 opset 下是否齐全两者取交集。如果模型里有较新的算子比如某些注意力相关的融合算子可能需要更高版本那就要推动下游升级或者在上游做等价替换。还有一个实操细节导出时如果某个算子在当前 opset 下没有实现工具通常会给出 warning 并把它降级成一个组合实现或者直接报错。warning 不能忽略它往往意味着图里多了一堆小算子推理性能会打折。我遇到这种情况会去看具体是哪个算子能换写法就换写法换不了再考虑升 opset。3.3 eval 模式、参数冻结与常量折叠model.eval()这件事前面提过这里展开说清楚它影响了什么。它影响的是所有训练态和推理态行为不同的层BatchNorm 用滑动统计量而不是 batch 统计量Dropout 直接透传一些自定义层如果有类似的开关也会切过去。导出前不切图里固化的就是训练态行为结果必然错。常量折叠是另一个容易被忽视的优化。导出时工具会把能提前算出来的子图直接算掉比如连续的 reshape、transpose、常量加法。这个开关默认是开的绝大多数情况应该保持开启它能让图更小、推理更快。但如果你的模型里有依赖运行时状态的常量计算折叠可能改变语义这时候需要关掉它并逐段核对。参数冻结指的是把参与训练但不参与推理的分支裁掉。比如一些模型带辅助分类头只在训练时用导出时应该让 forward 走推理分支。这件事没法靠参数解决只能改代码用if self.training之类的开关把分支隔离干净再配合 eval 模式让 trace 只走推理路径。4. 实操手把手导出一个可用的 ONNX 模型4.1 环境准备与最小可运行例子先把环境摆好。我一般用一个干净的虚拟环境装三个东西对应版本的 PyTorch、onnx、onnxruntime。前两个负责导出和结构检查第三个负责推理验证。版本按前面说的矩阵锁死别用 latest。下面是一个最小可运行例子结构简单但五脏俱全包含了卷积、归一化、激活、池化和全连接足够演示完整流程。import torch import torch.nn as nn class DemoNet(nn.Module): def __init__(self, num_classes10): super().__init__() self.conv nn.Conv2d(3, 32, kernel_size3, padding1) self.bn nn.BatchNorm2d(32) self.act nn.ReLU(inplaceTrue) self.pool nn.AdaptiveAvgPool2d(1) self.fc nn.Linear(32, num_classes) def forward(self, x): x self.conv(x) x self.bn(x) x self.act(x) x self.pool(x) x torch.flatten(x, 1) return self.fc(x) model DemoNet().eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, demo_net.onnx, input_names[image], output_names[logits], dynamic_axes{ image: {0: batch, 2: height, 3: width}, logits: {0: batch}, }, opset_version17, do_constant_foldingTrue, )几个点值得单独说。model.eval()在构造完就调用了不能省。dummy_input的 shape 决定了默认图结构虽然声明了动态轴但 trace 时走的还是这个具体 shape 的路径所以 dummy 的取值要贴近真实输入分布避免某些分支没被走到。opset_version显式指定不依赖默认值默认值会随框架版本变是隐性风险。4.2 关键参数逐个拆解input_names和output_names看着不起眼实际很重要。它们决定了图里输入输出节点的名字下游加载模型时就是靠这些名字来绑定数据的。不指定的话工具会生成input.1、output.1这种自动名可读性差而且不同版本生成的规则可能变。我见过上下游因为名字对不上互相扯皮的情况所以这两个参数我从来都是显式写。dynamic_axes的坑前面讲过这里补充一个实操技巧导出后用 onnx 工具把图的输入输出 shape 打出来确认动态维度确实变成了符号名而不是数字。这一步花不了一分钟能挡掉后面一堆问题。do_constant_folding建议保持 True。它的作用是把能提前算的常量运算折叠掉减小图规模。唯一需要关掉的情况是模型里有依赖运行时状态的常量计算但这种设计本身就不推荐遇到了应该改模型而不是关优化。opset_version的选择逻辑前面说过补充一点如果你不确定该用哪个可以先按下游支持的最高版本导一次跑通了再往下试找到一个稳定且够用的版本。不要一开始就追求最新稳定性优先。还有一个参数export_params默认 True表示把权重一起写进文件。除非你确实要做权重外置否则保持默认。外置权重会多一个文件部署时容易漏拷我一般不用。4.3 多输入、多输出与自定义算子多输入模型的导出把 dummy input 换成一个元组就行顺序和 forward 的参数顺序一致。input_names也按同样顺序列。这里最容易出错的是顺序对不上导出能成功推理时数据绑错位置结果全错还查不出来。dummy (torch.randn(1, 3, 224, 224), torch.randn(1, 128)) torch.onnx.export( model, dummy, multi_input.onnx, input_names[image, text_feat], output_names[logits, probs], dynamic_axes{image: {0: batch}, text_feat: {0: batch}}, )多输出同理forward 返回元组output_names按顺序对应。我建议输出名带上语义比如logits、probs、boxes别用out0、out1下游对接的时候能省很多沟通。自定义算子是最麻烦的一类。ONNX 标准里没有的算子导出会失败或者生成一个ATen域的节点很多推理引擎不认。处理思路有三条一是用标准算子组合出等价实现这是最干净的方案二是注册自定义算子需要下游引擎配合支持成本高三是导出后手工改图用等价子图替换适合一次性场景。我的优先级是能组合就组合实在不行再谈自定义。组合实现可能会带来一些性能损失但换来的是可移植性绝大多数场景这笔账是划算的。4.4 大模型与显存受限场景参数量大的模型导出时有个现实问题trace 过程需要把模型完整加载还要跑一次前向显存占用不小。如果单卡放不下可以用几个手段缓解。一是用 meta device 或者半精度加载导出前再转回需要的精度。二是分段导出再拼图适合结构清晰的模型但要处理跨段边界。三是用torch.onnx.export的device参数让导出在指定设备上完成新版支持这个能力。另外提醒一句大模型导出时间可能很长几分钟到几十分钟都有。导出过程别用交互式环境写成脚本跑日志重定向到文件出问题能回溯。5. 导出后验证怎么确认它没悄悄坏掉5.1 结构检查与图可视化导出成功不等于导出正确。第一步是做结构检查用 onnx 自带的检查器把图扫一遍能发现算子版本不匹配、节点连接错误、shape 推断失败这类问题。这一步很快应该写进脚本自动执行。import onnx model onnx.load(demo_net.onnx) onnx.checker.check_model(model) # 打印输入输出信息确认动态轴生效 for inp in model.graph.input: print(inp.name, [d.dim_param or d.dim_value for d in inp.type.tensor_type.shape.dim]) for out in model.graph.output: print(out.name, [d.dim_param or d.dim_value for d in out.type.tensor_type.shape.dim])打印出来的维度里动态轴应该显示成名字比如batch静态轴显示成数字。如果该动态的地方显示成了 1说明dynamic_axes没配对回去改。图可视化工具能把计算图渲染出来节点和连接一目了然。我一般用它做两件事确认关键算子确实存在、没有被降级成奇怪的组合确认图的输入输出节点名字符合预期。图太大时不用全看看首尾和中间几个关键节点就够。5.2 数值对齐用同一批输入跑两遍结构没问题只是及格线数值对不对才是关键。做法很直接准备一批真实分布的输入分别用原模型和 ONNX 模型跑一遍逐元素比对。import numpy as np import onnxruntime as ort sess ort.InferenceSession(demo_net.onnx, providers[CPUExecutionProvider]) input_name sess.get_inputs()[0].name x torch.randn(4, 3, 224, 224) with torch.no_grad(): ref model(x).numpy() onnx_out sess.run(None, {input_name: x.numpy()})[0] max_abs_err np.max(np.abs(ref - onnx_out)) print(max abs error:, max_abs_err)误差阈值要看任务。分类任务我一般看预测类别是否一致要求 100% 一致回归任务看最大绝对误差和相对误差阈值按业务容忍度定。如果误差超标先别急着改模型先排查是不是操作模式、输入 dtype、归一化处理哪里不一致。5.3 偏差来源排查表数值对不齐的原因有规律可循我把常见来源整理成一张表按出现频率排序遇到问题从上往下查。偏差来源典型表现排查手段忘记 eval 模式误差大且不稳定检查导出前是否调用 eval输入预处理不一致系统性偏移对比两边的归一化、通道顺序dtype 不一致小误差或类型报错确认 float32 对齐算子实现差异局部误差定位到具体层逐层对比动态轴配置错误shape 相关报错打印输入输出 shape常量折叠语义变化特定输入下出错关闭折叠重新导出对比逐层对比是个很有效的手段。把模型拆成若干段每段单独导出、单独比对能快速定位是哪一层开始出现偏差。这个方法费点时间但比盲猜高效得多。6. 常见问题与排查速查表6.1 典型报错与解法导出过程里的报错五花八门但高频的就那么几类。RuntimeError: Only tuples, lists and Variables are supported一般是 forward 返回值类型不对比如返回了字典或者自定义对象。解法是改成返回元组需要多输出的用元组组织。Unsupported operator算子不支持。解法是查这个算子在当前 opset 下的实现情况能换写法就换换不了就升 opset 或者做等价组合。TracerWarningtrace 过程中遇到了控制流或者动态行为图可能不完整。warning 不能无视要认真看是哪个操作触发的能静态化就静态化。Shape inference failedshape 推断失败通常是图里有动态维度没声明清楚。解法是补齐dynamic_axes或者检查是否有运行时才确定的 shape。6.2 动态 shape 踩坑实录动态 shape 是问题高发区值得单独说。最常见的错误是声明了动态轴但推理时喂了超范围的值比如声明了动态但实际实现里有隐含假设。有些模型内部用了 view 或者 reshape把 batch 维和空间维混在一起算batch 一变就崩。这种问题要在导出前用多个不同 shape 的输入测一遍别只测一个。第二个坑是动态轴声明了但下游没用上。有些引擎在加载时会根据图里的 shape 信息做静态优化如果它没正确识别符号维度优化后的图可能还是按固定 shape 跑。这种情况要在引擎侧确认它对动态维度的处理策略必要时在加载时显式指定。第三个坑是输出维度依赖输入。比如检测模型输出框的数量和输入分辨率相关这种动态性在 ONNX 里表达起来比较别扭通常的做法是固定最大输出数用无效值填充。这个方案不优雅但实用我见过不少项目这么做。6.3 部署端的隐性坑导出侧做对了部署侧照样能翻车。我遇到过几次问题都不在模型本身。一是输入布局不一致。PyTorch 用 NCHW有些推理引擎默认 NHWC加载时如果不做转换结果会错得离谱但不会报错。这种问题要靠对齐验证发现不能靠肉眼。二是预处理没对齐。训练时的归一化参数、通道顺序、resize 插值方式任何一个和部署侧不一致都会导致精度下降。我的做法是把预处理参数写进模型元信息部署侧照抄不做二次解释。三是线程数和内存配置。推理引擎的线程数、内存池大小、图优化级别都会影响性能和稳定性默认值不一定适合你的场景。上线前要做一轮压测把配置调到稳定区间。7. 工程化把导出固化成可复现的流水线7.1 目录结构与命名约定零散的导出脚本早晚会变成技术债。我的习惯是给每个模型建一个独立的导出目录里面固定放几样东西导出脚本、依赖清单、验证脚本、产物文件、元信息文件。命名带上模型名和版本号比如detector_v3_opset17.onnx一眼能看出是什么、什么版本、什么 opset。元信息文件里记录导出时间、框架版本、opset 版本、输入输出名和 shape、动态轴配置、验证误差、负责人员。这份文件的价值在半年后体现那时候没人记得当时怎么导的全靠它。7.2 性能优化与量化导出只是第一步后面还有优化空间。图优化层面推理引擎通常会自动做算子融合和内存复用你可以在加载时选择优化级别。量化层面如果业务允许精度损失把 float32 转成 int8 能带来显著的性能提升但需要校准数据集并且要做充分的精度回归。量化的坑在于校准集的选择。校准集分布和真实数据偏差大量化后的精度会掉得很厉害。我的做法是从真实流量里采样覆盖主要场景样本量按模型复杂度定一般几百到上千张足够。量化后必须重新跑一遍对齐验证阈值可以比 float32 放宽但不能不设。7.3 CI 里加一道回归关卡导出流程最怕的是改了模型忘了重导或者重导了但没验证。解法是把导出和验证写进持续集成模型代码一有变更就自动跑一遍验证不通过直接卡住。具体做法是准备一份小的验证数据集固定随机种子每次导出后自动比对输出。这个流程搭起来要花点时间但能挡掉绝大多数低级错误。我在几个项目里推过这套跑起来之后导出相关的线上问题基本归零。8. 一些踩坑之后才明白的事最后分享几条个人经验都是吃过亏换来的。dummy input 的取值比想象中重要。它不是随便填个形状就完事取值分布要贴近真实数据否则某些依赖数值的分支可能走到不同路径。我遇到过一次dummy 用全零导出后模型对某些输入输出异常换成随机正态就正常了。原因是模型里有基于输入统计量的逻辑全零输入走了特殊路径。导出脚本要能重复跑。不要用交互式环境一步步敲写成脚本参数化任何一次导出都能用同一条命令复现。这听起来是基本功但我见过太多人导出参数改了没记录出问题只能从头试。验证要自动化。手动比对一次两次可以长期靠人眼不现实。把验证写进脚本每次导出自动跑误差超标直接报错。这套机制建立起来之后你对导出质量的信心会完全不一样。还有一个心态上的建议导出失败别急着换方案。先看报错很多问题是有明确原因的改一行代码就能解决。我见过有人一报错就换框架、换工具链折腾一圈回到原点。导出这件事的复杂度大部分在细节不在选型。这个领域后续还能往深里做的是量化感知导出和跨版本兼容矩阵的自动化维护前者能进一步压推理成本后者能减少版本升级带来的返工。如果手头的模型对延迟敏感这两块值得投入。
阅读完成 · 觉得有帮助？
