模型训练完之后怎么落地是很多人卡住的地方。我见过不少团队训练脚本写得飞起指标也漂亮但一到部署环节就傻眼训练环境里跑得好好的网络换个运行时、换个硬件、换个语言栈输出就对不上。PyTorch 导出 ONNX这件事本质上就是解决这个落地断档的问题——它把动态图里那套 Python 运行时依赖剥掉压成一张静态的计算图交给通用推理引擎去跑。这篇内容面向的是已经能训模型、但还没把模型真正推上生产环境的同学也适合做算法工程、推理优化、端侧部署的朋友参考。我会从为什么选 ONNX 讲起把导出流程、算子踩坑、动态轴配置、精度对齐、性能验证这一整套东西掰开来说尽量让你照着做就能复现而不是停留在跑通 demo的层面。1. 先想清楚为什么要导出而不是直接扔个 pth 过去1.1 ONNX 到底解决了什么问题先说清楚 ONNX 是什么。它的全称是 Open Neural Network Exchange直白点讲就是一套跨框架的模型描述规范用一张由节点和边组成的有向图把网络的计算过程记录下来每个节点是一个算子每条边是张量的流动方向。它不关心你最初用的是哪个训练框架只关心输入什么、做什么运算、输出什么。这件事的价值在哪PyTorch 的默认模式是动态图eager mode模型的前向逻辑是 Python 代码一行一行解释执行的。这意味着你的推理服务必须拖着整个 Python 解释器、PyTorch 运行时、甚至 CUDA 上下文一起走。部署到服务端还好一旦要往端侧、嵌入式、或者别的语言栈C、C#、Java里塞这套依赖就是个包袱。而 ONNX 把计算过程固化成图之后推理引擎可以用 C 实现、可以做图优化、可以做算子融合、可以针对特定硬件做编译这是纯 PyTorch 脚本很难做到的。我通常把这件事类比成把一份手写的菜谱翻译成标准化的工序卡。手写菜谱PyTorch 模型只有原作者看得懂换个厨师推理引擎就得重新理解工序卡ONNX规定了每一步用什么原料、做什么动作、产出什么半成品任何合格的厨房都能照着做。1.2 导出前必须搞清楚的三个前提在动手敲torch.onnx.export之前有三件事必须先确认否则后面全是返工。第一模型必须处于 eval 模式。这一点被无数人忽略。model.train()状态下BatchNorm 用的是当前 batch 的统计量Dropout 会随机置零导出的图里这些行为会被固化或者产生不确定结果。正确做法是先model.eval()再用torch.no_grad()包住导出过程。我踩过一次坑某个带 Dropout 的分类模型导出时忘了切 eval验证集精度掉了七八个点查了半天以为是导出算子有问题其实是推理时还在随机丢弃。第二输入的形状要想明白。ONNX 导出的图是静态的除非你显式声明动态轴。如果你的服务端会接受不同 batch size 或者不同分辨率的图片那导出时必须把对应的维度标成 dynamic否则引擎会按导出时的固定 shape 做优化喂进去不一样大小的输入直接报错。第三算子兼容性要提前探底。PyTorch 有大量算子在 ONNX 标准里没有对应实现或者只有高版本 opset 才支持。比较典型的是各种自定义算子、某些特殊的下标索引、复杂的 reshape 组合。导出脚本会给你警告但警告往往被日志淹没。我的习惯是导出后立刻用 onnx 库做一次 checker再看算子列表里有没有可疑项。1.3 版本矩阵这件事别等出问题才查ONNX 导出对版本极其敏感。PyTorch 版本、onnx 库版本、onnxruntime 版本、opset 版本四者之间有一条隐性的兼容链。opset 是 ONNX 的算子集版本PyTorch 每个版本支持的 opset 上限不同而 onnxruntime 又只认某个范围内的 opset。下面这张表是我在几个常见组合上实测的记录供参考PyTorch 版本建议 opsetonnx 库版本onnxruntime 版本备注1.13.x171.141.14稳定组合动态轴支持好2.0.x171.141.15注意 torch.compile 与导出不要混用2.1.x17~181.151.16部分 attention 算子需 182.2.x181.161.17推荐固定 opset 显式声明注意不要盲目追最新 opset。有些推理引擎对最新 opset 的支持滞后导出成功不代表能加载成功。稳妥策略是选一个你的训练框架和推理引擎都明确支持的中间版本并在项目里写死。2. 核心导出流程与关键参数拆解2.1 最小可跑通的导出脚本先把骨架搭起来。下面这段是最小可用版本我会在注释里说明每个参数的作用import torch import torch.nn as nn class DemoNet(nn.Module): def __init__(self): super().__init__() self.conv nn.Conv2d(3, 16, 3, padding1) self.bn nn.BatchNorm2d(16) self.act nn.ReLU() self.pool nn.AdaptiveAvgPool2d(1) self.fc nn.Linear(16, 10) 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() model.eval() # 必须切 eval dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, demo.onnx, export_paramsTrue, # 权重一并写入 opset_version17, # 显式指定别用默认 do_constant_foldingTrue, # 常量折叠优化 input_names[input], output_names[logits], dynamic_axes{ input: {0: batch_size, 2: height, 3: width}, logits: {0: batch_size}, }, verboseFalse, )几个参数值得单独讲。export_paramsTrue表示把训练好的权重作为常量嵌进图里这是常规做法只有极少数需要权重外置的场景才设 False。do_constant_foldingTrue会让导出器把能静态算出来的子图提前算掉比如连续的 reshape 和常量相加能显著减小图体积。dynamic_axes是最容易出错的地方后面单独一节说。dummy_input这个参数名字起得很诚实——它只是用来走一遍前向、记录计算轨迹的。但它的形状和 dtype 会影响导出的图所以一定要跟真实推理时的输入对齐。我见过有人用torch.randn(1,3,224,224)导出线上喂的是float16的输入结果一加载就报 dtype 不匹配。2.2 dynamic_axes 的正确写法与隐性代价动态轴是新手最容易翻车的地方。先说结论只把真正会变的维度标成 dynamic。标多了会带来两个后果一是图优化受限很多融合只有静态 shape 下才能做二是某些推理引擎在动态 shape 下会走更保守的 kernel性能反而下降。写法上有个细节dynamic_axes的 key 必须是input_names或output_names里出现过的名字不能直接写模型里的变量名。上面例子里输入叫input输出叫logits所以 key 用这两个。还有一个隐蔽的坑如果模型里有view、reshape、flatten这类操作动态 batch 可能会让 shape 推断失败。原因是这些算子需要知道具体维度值。torch.flatten(x, 1)在动态 batch 下通常没问题因为它只合并后几维但如果你写的是x.view(batch, -1)且后续又依赖batch的具体数值就可能出问题。比较稳的写法是用torch.flatten替代手写 view或者在 forward 里避免对 batch 维度做算术。我曾经处理过一个分割模型输入动态、输出也动态导出后 onnxruntime 加载报错说某个 Resize 节点的 scale 是动态的但 axes 不对。排查下来是模型里用了nn.Upsample(scale_factor2)动态输入下 scale 变成了运行时计算的张量。解决办法是把scale_factor换成固定的size参数或者干脆把输入尺寸固定住。这件事告诉我不是所有模型都适合全动态导出很多时候固定分辨率反而更省心。2.3 用 torch.jit.trace 还是用 dynamo 导出PyTorch 的 ONNX 导出现在有两条路。老路是基于 TorchScript 的 trace也就是上面那个torch.onnx.export的默认行为新路是 PyTorch 2.1 之后引入的 dynamo 导出器参数是dynamoTrue。两者的区别值得说清楚。trace 的思路是喂一个样例跑一遍把执行的算子记录下来所以它对控制流极其敏感——if分支如果依赖输入数据trace 只会记录样例走的那条路另一条分支就丢了。而 dynamo 的思路更接近编译它会尝试追踪 Python 层的控制流并转成图覆盖率高一些。但 dynamo 导出并不是万能良药它对某些自定义算子和复杂 Python 逻辑的支持还在完善中而且导出失败时的报错信息往往不如 trace 直观。我的实践结论是先用 trace导出失败或者发现算子丢失再试 dynamo。trace 这条路成熟、社区案例多、调试工具齐全能覆盖绝大多数 CNN 和常见 Transformer 结构。如果你确实有数据相关的控制流更稳妥的做法是把它改写成torch.where之类的张量运算让图变得可 trace而不是强行依赖导出器去猜。3. 导出之后的验证不验证等于没导出3.1 onnx.checker 与图结构自检导出成功不代表模型正确。第一件事是跑一次结构检查import onnx model_onnx onnx.load(demo.onnx) onnx.checker.check_model(model_onnx) print(graph structure ok) # 看算子清单 from collections import Counter op_types Counter(node.op_type for node in model_onnx.graph.node) for k, v in op_types.most_common(): print(f{k}: {v})checker 只能保证图结构合法不能保证行为正确。真正有价值的是算子清单——如果里面出现了ATen、prim::开头的东西说明有算子没被正确翻译退化成了 PyTorch 自定义算子这类图在通用推理引擎里根本跑不了。我一般会扫一遍清单确认没有这类漏网之鱼同时看看有没有意料之外的If、Loop节点。另外一个检查点是输入输出的 shape 和 dtype。model_onnx.graph.input里能看到导出器记录的信息如果发现某个维度变成了unk__0这种符号名说明动态轴生效了这是正常的但如果你本来想要静态 shape 却出现了符号维那就是配置问题。3.2 数值对齐这一步绝不能省结构检查过了接下来必须做数值对齐。思路很朴素同一份输入分别喂给 PyTorch 模型和 onnxruntime比较两边输出的最大绝对误差。import numpy as np import onnxruntime as ort x torch.randn(1, 3, 224, 224) with torch.no_grad(): torch_out model(x).numpy() sess ort.InferenceSession(demo.onnx, providers[CPUExecutionProvider]) onnx_out sess.run(None, {input: x.numpy()})[0] diff np.abs(torch_out - onnx_out) print(max abs diff:, diff.max()) print(mean abs diff:, diff.mean())判据怎么定我的经验是CPU 上对同一份 float32 输入最大绝对误差在1e-5量级以内算正常1e-3量级就要警惕超过1e-2基本可以判定有问题。这个阈值不是死的跟模型深度有关——层数越深误差累积越明显几十层的网络出现1e-4也能接受。注意如果原始输出本身数值很大比如未归一化的 logits 有几十上百要用相对误差辅助判断别只看绝对误差。误差偏大时怎么查我的套路是先定位到第一个对不上的节点。做法是把 ONNX 图按节点切开用 onnxruntime 逐层跑和 PyTorch 里对应层的中间输出比。工具上可以用onnx.utils.extract_model把子图抽出来或者更省事的办法是在 PyTorch forward 里插 hook 记录每层输出然后对照。通常第一个出问题的节点就是元凶常见于归一化层、激活函数、以及各种带epsilon的算子。3.3 动态轴下的批量验证不能只测一条如果导出时开了动态轴验证就必须覆盖多个 shape。只测 batch1 是不够的你得测 batch1、2、8测不同分辨率测边界情况比如 batch 小于模型内部某个分组数。我遇到过一次事故模型里有 GroupNorm分组数是 8导出时 batch 固定为 8线上来了一条 batch1 的请求直接崩了——因为 1 不能被 8 整除。这类问题的排查成本很高因为测试环境全过、线上偶发。稳妥做法是写一个小脚本把可能的 shape 组合枚举一遍全部跑通再上线。这也是我为什么强调固定 shape 有时候更省心——如果你的业务场景输入尺寸本来就固定干脆别开动态轴省掉一整类问题。4. 高频踩坑与排查实录4.1 算子不支持与自定义算子怎么办这是导出环节最大的拦路虎。PyTorch 生态里有大量算子在 ONNX 标准里没有对等物比如某些特殊的插值模式、可变形卷积、复杂的注意力变体。现象通常有两种一是导出直接抛异常二是导出成功但图里出现com.microsoft或ATen域的自定义节点。处理思路分三层。第一层是换实现很多算子的等价写法是标准算子组合。比如F.interpolate的某些模式可以拆成 Resize 特定插值某些自定义的归一化可以用标准算子手工拼出来。第二层是降级如果某个算子只在训练时用推理时根本不需要那就用torch.onnx的算子屏蔽机制或者干脆在导出前把相关分支剪掉。第三层才是自定义算子注册给推理引擎写扩展算子这需要改引擎侧代码成本最高一般只有前两层走不通才考虑。我个人的优先级很明确能换实现就换实现能剪枝就剪枝实在不行才动自定义算子。因为自定义算子意味着你的 ONNX 文件不再是通用可移植的换一个推理引擎就得重写一遍失去了导出的最大意义。4.2 精度漂移的常见源头清单数值对不上是第二大高频问题。我把踩过的源头整理成下面这张表按出现频率排序现象可能原因排查手段解决方式整体误差 1e-3 量级常量折叠改变了浮点运算顺序关掉 constant_folding 重导比对接受误差或调整阈值某几层后误差突增BatchNorm 的 epsilon 或 momentum 参数不一致检查 bn 层参数是否被正确导出统一 eps 设置输出整体偏移输入预处理在两边不一致归一化、通道顺序用同一份原始输入跑预处理把预处理固化进图或统一到引擎侧只有部分样本错动态轴下 shape 推断出错打印中间张量 shape 对比固定相关维度全零或 NaNdtype 不匹配fp16 输入喂 fp32 图检查 input dtype导出时指定正确的 dtype这张表里我要特别强调预处理不一致这一条。大量所谓的导出精度问题根本不是导出器的问题而是训练时的预处理和服务端的预处理没对齐——比如训练用 PIL 读图、服务端用 OpenCV 读图BGR 和 RGB 反了或者归一化的 mean/std 在两边写法不同。这类问题的典型特征是误差极大且呈现规律性偏移用一张已知答案的图片测一下就能立刻确认。4.3 大模型导出的内存与体积问题模型越大导出越容易撞内存墙。torch.onnx.export默认会把所有权重作为常量塞进 protobuf几十亿参数的模型导出来的 onnx 文件能到几十 GBprotobuf 有单文件 2GB 的硬限制超了直接失败。应对手段有几个。一是启用外部数据格式让权重存到单独的.data文件里import onnx from onnx.external_data_helper import convert_model_to_external_data model_onnx onnx.load(big.onnx) convert_model_to_external_data( model_onnx, all_tensors_to_one_fileTrue, locationbig.onnx.data, size_threshold1024, convert_attributeFalse, ) onnx.save_model(model_onnx, big_external.onnx)二是导出时就把大模型切成子图分块导出尤其是 Transformer 这类结构清晰、可以按层切的模型。三是考虑用 fp16 或者更激进的量化来压体积但这会影响精度得做完整的数值验证。第三个是导出过程本身的内存开销。trace 模式下导出器要把样例前向的所有中间张量记录下来显存/内存峰值可能比正常推理高好几倍。我的做法是把 dummy batch 设成 1并且用torch.no_grad()包住能省下不少。5. 导出后的性能与稳定性验证5.1 别只看能不能跑要看跑得快不快导出成功、数值对齐很多人就收工了。但真实项目里还有一关性能。ONNX onnxruntime 不一定比原生的 PyTorch 快这取决于图优化是否生效、算子融合是否命中、以及有没有走对执行提供器Execution Provider。最基础的测法是同一份输入跑多次取平均对比 PyTorch 和 onnxruntime 的耗时import time def bench(fn, x, warmup10, iters50): for _ in range(warmup): fn(x) t0 time.perf_counter() for _ in range(iters): fn(x) return (time.perf_counter() - t0) / iters with torch.no_grad(): t_torch bench(lambda v: model(v), torch.randn(1,3,224,224)) print(pytorch:, t_torch) x_np np.random.randn(1,3,224,224).astype(np.float32) t_onnx bench(lambda v: sess.run(None, {input: v}), x_np) print(onnxruntime:, t_onnx)注意一定要有 warmup。onnxruntime 第一次跑会做图优化、内存分配、算子选择耗时是稳态的好几倍不做 warmup 测出来的数没意义。如果发现 onnxruntime 反而更慢先查三件事一是 provider 用的是不是 CPU 默认实现有没有开对应硬件的加速二是图优化等级是否拉满sess_options.graph_optimization_level三是动态轴是不是开太多了导致优化器没法做静态 shape 的融合。5.2 图优化等级与执行提供器的选择onnxruntime 的SessionOptions里有几个关键旋钮。graph_optimization_level有四个档位从禁用一直到全部优化默认是全部。生产环境一般保持默认但调试时可以把优化关掉用来判断某个精度问题是不是优化引起的。执行提供器是另一个重点。CPU 上默认是 CPUExecutionProvider如果机器有对应加速硬件要显式指定对应的 provider并注意 provider 的注册顺序——onnxruntime 会按顺序尝试把最快的放前面。注册顺序写反了性能可能差好几倍。还有一点容易被忽略动态轴带来的 shape 变化会让 onnxruntime 反复做 shape 推断和内存重分配。如果你的服务输入 shape 种类有限可以考虑做几个固定 shape 的 session 分别加载用的时候按 shape 路由。这个策略在高并发场景下收益很明显。5.3 上线前的检查清单我现在每次交付一个导出的 ONNX 模型都会过一遍下面这个清单模型是否 eval 模式导出Dropout 和 BN 行为是否确认无误输入输出的 shape 和 dtype 是否与线上契约一致动态轴是否只标了必要的维度onnx.checker 是否通过算子清单里有无 ATen / prim 残留数值对齐误差是否在阈值内是否覆盖了单条和批量两种输入动态轴场景下是否枚举测试了足够多的 shape 组合是否测过冷启动加载时间大模型是否用了外部权重性能是否对比过基线图优化和执行提供器是否配置正确ONNX 文件是否做了版本归档与训练产物的对应关系是否记录这几条看起来啰嗦但每一条背后都是一次真实的返工。尤其是最后一条——导出的 ONNX 文件本质上是一个构建产物必须和训练时的模型权重、代码版本绑在一起存档否则线上出了问题根本没法复现。我用过的最省事的办法是给文件名带上版本号和后缀同时把导出脚本、参数配置、依赖版本一起提交到版本库让重新导出一次成为一条可执行命令。最后分享一个我最近养成的小习惯导出脚本里加一个--verify开关默认开启。脚本跑完导出后自动执行 checker、数值对齐、shape 枚举、以及一次简短性能测试任何一项不过就返回非零退出码。这个做法把上面那堆检查固化成了流程比靠记忆力可靠得多。踩过的坑不想再踩第二遍能交给脚本的就别交给人。