模型导出这件事看起来是整个流程里最不起眼的一步可它偏偏是我见过翻车最多的地方。很多同学在 PyTorch 里把模型训得风生水起一到导出 ONNX 就开始原地踩坑动态轴写错、算子不兼容、导出结果和原模型对不上甚至连最后一个维度顺序都能让推理直接起飞。这个阶段要是没处理好后面量化、裁剪、端侧推理全都会被死死拖住。我这次想聊的是项目推进到 Phase A 第三步时做的一次完整模型导出。目标很明确把手里的 PyTorch 模型稳定地转成 ONNX为后续的 INT8 量化、边缘设备部署和跨平台推理铺路。这篇文章会从“为什么要把导出放在这个阶段”讲起然后一步步拆解 PyTorch 转 ONNX 的完整流程、YOLO 系列导出的典型坑、如何用 ONNX Runtime 做一致性校验、量化以及与端侧格式如 KMODEL 的转换链路最后补一段 TTS 场景里导出 ONNX 的特殊考虑。适合正在做模型部署、推理优化、边缘计算或者刚被 ONNX 导出折腾到怀疑人生的朋友。1. 为什么导出选 ONNX从训练态到部署态的关键跨越1.1 导出不只是格式转换更是一次“模型固化”训练阶段和部署阶段的模型本质上是两种状态。训练时我们看重的是梯度回传、动态图调试、参数更新所以 PyTorch 模型里到处都是灵活的 Tensor 操作你随时可以打断、打印、改代码。但部署阶段不需要梯度不需要 autograd 图也不需要 Python 环境每秒都在解释执行。部署只需要一件事给定输入稳定地、快速地、在目标硬件上算出输出。所以“导出 ONNX”看起来是在换格式实际是在做一次模型固化。模型的结构、权重、算子顺序、输入输出接口全都被写进一张静态计算图里。图一旦生成就不再依赖你原来那套 nn.Module 定义也不再依赖训练代码里各种动态控制流。这个过程越早意识到后面遇到的幺蛾子越少。放在 Phase A 第三步来做还有一个很现实的原因模型还没进入大规模调参和复杂分支开发的时候结构相对干净。如果等到模型里堆了十几个自定义算子、各种动态 shape 分支再导出光是修算子映射就能耗掉一两天。早一步导出等于提前把“能不能部署”这件事的风险暴露出来而不是等到项目快上线才发现模型根本转不动。1.2 ONNX 在模型交换链路里的角色ONNX 全称 Open Neural Network Exchange简单说就是模型界的“普通话”。它不绑定某个深度学习框架也不依赖某个硬件厂商而是一个开放、可读、有明确版本规范的计算图描述格式。很多框架都能导出或导入 ONNX这让我们在模型分发的时候不需要强迫对方也装一套 PyTorch 环境。这种中间格式的思路在其他领域也很常见。用过 Blender 里 better fbx importer exporter 这类插件的人应该能秒懂FBX 之于三维模型交换就像 ONNX 之于神经网络模型交换。它解决的不是“让你的模型跑得最快”而是“让你的模型能进到别人的工具链里”。ONNX 生态里还有很多配套工具比如 ONNX Runtime 可以直接加载推理onnx 库可以检查图结构量化工具可以做 INT8 压缩各家边缘 AI 芯片厂商也往往会提供从 ONNX 到私有格式的转换入口。选择 ONNX意味着你没有被绑死在单一框架或单一硬件上。1.3 导出时机的选型风险导出太晚模型复杂度会拖累转换导出太早又可能模型结构还没定型每次改动都要重新导。我的经验是把导出节点放在“模型结构基本稳定但超参还在微调”的阶段。在这个阶段先跑通一条完整的“训练 → 导出 → 验证 → 量化 → 部署”链路哪怕先用一个简化版模型验证也行。链路通了后面每次模型更新都只是重复执行固定流程心里会非常踏实。也要提前接受一个事实静态图导出不是万能的。PyTorch 的动态控制流如果太野比如循环次数随输入长度变化、分支里临时创建中间张量ONNX 导出时很容易失败或导出一个“看似能用实际行为不对”的图。所以在导之前最好对模型内部做一次“减负”能用固定 shape 就固定能去掉的 Python 分支就去掉。2. 导出 ONNX 之前的模型检查清单2.1 权重状态与 eval 模式这是最基础、也最容易被忽略的一点。导出前一定要确认模型处于 eval 模式并且用 torch.no_grad() 包住导出过程。原因不复杂如果模型还停留在 train 模式BatchNorm 层还在用当前 batch 的统计量更新 running_mean 和 running_varDropout 层还在随机丢弃神经元你导出来的计算图虽然结构没错但行为已经和你想部署的模型不一致了。另一个容易踩的坑是权重版本对不上。很多人会加载训练过程中保存的 checkpoint但 checkpoint 里可能同时有 model_state_dict 和 optimizer_state_dict。如果直接 torch.load 之后没取对键名模型虽然能 forward但用的可能是随机初始化的权重。导出来之后精度自然完全不对。所以在导出脚本里我的习惯是先用一个独立的小脚本加载权重在一批固定输入上打印几层输出的均值和方差确认权重真的加载对了再开始导出。2.2 确定输入输出的张量形状和类型ONNX 导出需要提供 dummy inputPyTorch 会根据这个输进去的 Tensor 实际走一遍 forward把计算图记录下来。因此 dummy input 的形状、数值类型、数据排布直接决定导出图的输入约定。很多人在这一步图省事用 torch.randn(1, 3, 224, 224) 随便试一下结果模型内部要求的是 NHWC 或者归一化到 0-1 范围的输入导出时又没显式说明部署阶段就开始乱套。这里有个习惯特别值得养成在导出脚本里把输入输出命名做好比如 input 和 output如果有多个输出就逐个命名。命名不是写给自己看的而是给后面的 ONNX Runtime、TensorRT、RKNN、地平线工具链等下游环节看的。很多工具链默认把输入叫作 data输出叫作 prob但如果你在导出时定了更明确的语义名后续接后处理时要轻松很多。2.3 固定尺寸还是动态尺寸先算清推理场景这是导出 ONNX 时最灵魂的决策之一。dynamic_axes 参数允许你把 batch、height、width 设成动态也就是模型可以接收不同尺寸的输入。听起来很美好但代价是计算图的很多优化没法做某些硬件加速器根本不支持动态 shape甚至某些算子在后端会被迫 fallback 到很慢的实现。我一般先问自己一个问题这个模型上线后输入尺寸真的会变吗如果模型输入永远是固定分辨率比如图像固定 640x640那就不加动态轴让整个图尽量“死”一点换取更稳的部署表现。如果确实需要支持多分辨率比如检测模型要适配不同长宽比图片那可以把 height 和 width 设成动态但 batch 可以固定成 1。这种“部分动态”的折中方案在工程上的可维护性要远好于把所有维度都放开。3. PyTorch 转 ONNX 实操最小示例到 YOLO3.1 最小可运行导出示例与参数解释先给一个最小例子把主流程跑通再慢慢往里面加细节。import torch import torch.nn as nn class SimpleModel(nn.Module): def __init__(self): super().__init__() self.conv nn.Conv2d(3, 16, kernel_size3, padding1) self.fc nn.Linear(16 * 32 * 32, 10) def forward(self, x): x torch.relu(self.conv(x)) x x.view(x.size(0), -1) return self.fc(x) model SimpleModel().eval() dummy_input torch.randn(1, 3, 32, 32) torch.onnx.export( model, dummy_input, simple_model.onnx, input_names[input], output_names[output], opset_version17, dynamic_axes{input: {0: batch}, output: {0: batch}}, )这里每个参数都有实际含义。dummy_input 是触发 trace 的种子输入模型会顺着它把 forward 里的算子一个一个记录到图里。input_names 和 output_names 是给计算图的接口起名字后续推理时传 Tensor 就用这些名字匹配。opset_version 决定用哪个版本的 ONNX 算子集来映射 PyTorch 算子版本太低有些算子映射不出来版本太高又可能下游不支持我通常先选 17如果转化链路有要求再降。dynamic_axes 的写法值得多看两眼。{0: batch} 的意思是把第 0 维标记为动态逻辑名是 batch。这样导出的图在 batch 维度可以变化但其他维度还是固定。如果你希望 height 和 width 也动态就写成 {2: height, 3: width}。注意这里不能只写动态轴的逻辑名还要确保 forward 里对 shape 的处理不依赖于硬编码的常量否则 trace 的时候会把具体数字拍死。3.2 YOLO 系列导出 ONNX 的常见玩法YOLO 系列现在应该是检测领域导出 ONNX 最频繁的模型之一。网上随手一搜就是“yolo导出onnx模型”的教程但真正在项目里跑过一遍的人都知道坑不在 torch.onnx.export 这个函数本身而在模型结构的后处理部分。YOLOv5、YOLOv8 这类模型一般在导出前都要决定一件事NMS非极大值抑制到底放在 ONNX 里还是放在 ONNX 外面。大多数官方导出脚本默认不包含 NMS也就是模型只输出原始预测的坐标、置信度和类别概率后处理留给推理框架或者业务代码处理。这样做的好处是灵活你可以针对不同场景写不同的 NMS 策略坏处是部署方的代码要更复杂。如果想把 NMS 也放进 ONNX 图里PyTorch 里通常需要把 NMS 写成基于 torch.ops 的函数导出时再映射成 ONNX 的非标准算子。这会导致 ONNX 对 runtime 有特殊要求比如必须使用带自定义算子的 ONNX Runtime。我的建议是除非你的部署环境非常固定否则尽量导出不带 NMS 的版本先用 ONNX Runtime 跑通原始输出再在业务侧做后处理。这样排查问题的时候至少能把“模型计算错”和“后处理错”这两个环节切开。还有一个 YOLO 特有的问题导出时输入尺寸经常被自动改成固定值。很多官方仓库会读取一个 yaml 里的 imgsz 参数或者从训练配置里继承。你如果换了部署端要求的尺寸比如从 640 改成 320最好生成一个新的 ONNX 文件再做测试不要直接拿旧文件硬凑。3.3 用 ONNX Runtime 做一致性校验导出成功不等于导出正确。ONNX 文件是不是真的等价于原来的 PyTorch 模型必须用推理结果来验证。这里我固定用一套小脚本import numpy as np import onnxruntime as ort import torch pt_output model(dummy_input) sess ort.InferenceSession( simple_model.onnx, providers[CPUExecutionProvider], ) ort_input_name sess.get_inputs()[0].name ort_output_name sess.get_outputs()[0].name ort_output sess.run( [ort_output_name], {ort_input_name: dummy_input.numpy()}, )[0] np.testing.assert_allclose( pt_output.detach().numpy(), ort_output, rtol1e-03, atol1e-05, )校验时注意两个地方。第一PyTorch 的输出和 ONNX Runtime 的输出必须来自同一个输入且输入的数据处理方式完全一致。第二float32 模型经过 ONNX Runtime 计算后某些算子可能有微小的数值误差所以一般用 np.testing.assert_allclose 设置 rtol 和 atol而不是要求完全相等。如果误差超过 1e-2那基本可以判定导出过程有 bug不要强行用“精度损失”来解释。这个一致性校验脚本每次模型结构改完都要重新跑。我一般会把它写成一个 pytest 用例纳入 CI 流程里哪怕不跑完整训练也可以定期验证导出链路没有因为框架版本升级而悄悄失效。4. 导出之后的下游链路INT8 量化与边缘转换4.1 用 ONNX Runtime 给模型做 INT8 动态量化导出 ONNX 不是为了留在电脑上看而是要往部署走。部署到服务器 GPU 上可以用 FP16部署到边缘设备和端侧容量和带宽往往比较紧张这时候 INT8 量化就是最常用的压缩手段。ONNX Runtime 提供了现成的量化工具最简单的量化方式是动态量化from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( simple_model.onnx, simple_model_int8.onnx, weight_typeQuantType.QInt8, )动态量化只把权重量化成 INT8激活值在推理时动态决定缩放因子不需要校准数据适合快速试水。但动态量化对推理速度的提升有限真正要压榨性能得做静态量化。静态量化需要准备一批有代表性的校准数据让工具统计每一层激活值的分布再确定合适的缩放因子。做 INT8 量化时我最深的体会是不要只盯着压缩率。量化完一定要重新跑一遍业务测试集观察精度下降是否在可接受范围。有的模型对量化非常敏感可能一个分支结构稍复杂INT8 就直接掉两三个点。这时候可以尝试混合量化把敏感层保留为 FP16 或 FP32其余层用 INT8。ONNX Runtime 的 quantization 接口里也有 methods 和 nodes_to_exclude 相关配置可玩性相当高。4.2 从 ONNX 到 KMODEL 这类端侧格式的注意事项如果说 ONNX 是“人人都能读的通用语言”那 KMODEL 就是某一类端侧 AI 加速芯片的“外语”。当前很多边缘 AI 芯片厂商都提供了从 ONNX 导入模型的工具链这些工具链会把 ONNX 计算图做算子映射、内存规划、指令生成最终产出一个只能在自家硬件上运行的模型文件。从 ONNX 转 KMODEL 的路上踩得最多的坑就是算子兼容。ONNX 的算子集很大但端侧芯片的 NPU 加速库覆盖面有限。像 torch.onnx.export 默认输出的某些算子比如比较好的动态 reshape、某些特殊的插值方式在转 KMODEL 时可能压根不支持要么报错要么被放到 CPU 上慢速执行。我的建议是在选模型结构时就先看算子白名单避开那些明显不适合端侧的结构比如过于复杂的 attention block 或者一大堆自定义 op。另外输入输出的形状约定在转换环节里会被再次放大。有些边缘工具链要求输入是 NHWC有些要求 NCHW还有些允许你通过配置文件指定。这个环节如果没对齐转换能成功但跑起来就是结果错。最稳妥的做法是遵循“先 ONNX Runtime 确认 → 再量化 → 再转 KMODEL → 再硬件上校验”的链路每一步都单独验收不要一口气跳到最后。4.3 TTS 场景下导出 ONNX 的额外功课TTS 模型这些年也越来越依赖 ONNX 做发布。你去看开源社区里的 TTS 推理引擎比如 sherpa-onnx它就是用 ONNX 作为主要模型格式来分发和推理。好处很明显模型作者不需要为每个平台重新训练用户只要把 ONNX 文件放进引擎就能跑还支持移动端和桌面端。但 TTS 模型导出 ONNX 和图像模型很不一样。首先是可变长度问题同一句话不同时长梅尔谱帧数不一样声学模型和声码器经常需要处理动态时序长度。其次是流式推理很多实时语音合成系统会把模型切块每块只能看到局部输入同时还要维护一个状态缓存比如 LSTM 或 GRU 的 hidden state。如果你在 PyTorch 里写的是带 state 的递归结构导出 ONNX 时就不仅要导出主运算图还要把 state 也作为显式输入输出。我自己的建议是TTS 模型导出时不要急着把整个声学模型和声码器拼在一起导成一个大 ONNX。先各自导出、各自验证确认单模块的输入输出和中间特征匹配得上再考虑合成一个推理 pipeline。否则模型一旦变长中途某一个 feature 统计量偏了你根本定位不到是在声学模型还是声码器里出了问题。5. ONNX 导出常见问题与排查实录5.1 动态轴导致 graph 形状推断失败这是高频报错常见信息类似 “Cannot infer shape” 或 “Shape inference failed”。多半是因为你在计算图里对张量做了 view、reshape、或者 shape 相关的 Python 分支导出的 ONNX 在动态 shape 下没法静态推断出后续某个算子的输入维度。遇到这种情况我建议分三步排查。第一步把 dynamic_axes 全部去掉改成固定尺寸看导出是否成功。如果固定尺寸下正常说明问题确实出在动态轴如果固定尺寸下也报错那问题在算子本身。第二步检查 forward 里是不是用了 x.size() 之后的 int 参与 Python 运算比如 if x.shape[2] 10这种写法在 trace 时会被当作常数处理动态 shape 下语义会错。第三步实在找不到原因就把出问题的部分拆成一个子模块单独导出定位。5.2 算子不兼容换算子还是自研扩展PyTorch 算子数量远远超过 ONNX 标准算子所以导出映射失败非常正常。常见信息比如 “Unsupported opset or missing ONNX op for xxx”。这时候我会先做三件事。先确认算子在 ONNX 里有没有替代品。比如某些高级的 index_put、scatter 操作可以改写成 ONNX 支持的 Gather、ScatterElements。第二个办法是升级或降低 opset_version有些算子在高版本 opset 里才有标准映射有些反向不兼容反倒低版本更保险。这里有个经验如果你要转端侧 KMODEL通常建议 opset 别用太新11 到 13 往往是最稳的因为各家工具链对新增算子的支持速度没那么快。如果 ONNX 标准算子真的覆盖不了那就只能走自定义算子路线。可以在导出时用 torch.onnx.register_custom_op_symbolic 把自定义算子映射到自定义命名空间然后在 ONNX Runtime 侧通过 CustomOp 注册实现。这个方案能解决问题但维护成本不低能避开就避开。5.3 导出后精度对不上从哪几步入手查导出后模型精度差是最让人头疼的问题。因为整个链路里任何一个环节出错最终都会表现为输出不正确。我会按下面这个顺序排查第一确认 PyTorch 那边是 eval 模式且没有梯度。第二对比同一个输入在 PyTorch 和 ONNX Runtime 里的逐层输出快速定位是哪个算子开始偏离。ONNX Runtime 的 IO Binding 和模型导出时的 intermediate outputs 可以帮助你做到这一点哪怕做不到逐层也可以在模型里临时加几个输出节点做对比。第三也是最容易被坑的就是数据预处理不一致。PyTorch 训练时如果你用了 ToTensor、Normalize、通道变换部署端用 ONNX Runtime 推理时也要完全复刻这套流程。很多人训练时用 RGB、归一化到 0-1但部署脚本读图后直接转成 BGR 或者忘了除以 255输出自然千差万别。这类问题不是 ONNX 的锅却往往被记到 ONNX 头上。6. 想让导出少踩坑我给自己定的几个默认动作写到最后分享几个我长期坚持的小习惯。每次导出 ONNX 之前我都会先拿一个小输入样例把模型从 checkpoint 加载、导出、ONNX Runtime 推理、结果比对这四步跑通然后再拿真实尺寸的数据去走完整流程。这个习惯帮我省了大量试错时间因为小样例跑的路径和大模型完全一样但报错信息更容易定位。另外我习惯把每一版导出的 ONNX 文件连同 opset 版本、动态轴配置、备注一起存好。这个模型可能是从 YOLO 导出的检测模型也可能是一个 TTS 的声学模型如果不过一段时间再回来看很可能忘了当初那版模型究竟是怎么导的。一条清晰的导出日志比任何口头文档都靠谱。还有一个不算技术但特别管用的经验模型导出这件事千万不要在最后一天集中做。无论是 PyTorch 转 ONNX、INT8 量化还是转 KMODEL每一个环节都有自己独立的风险。提前打通链路平时模型每更新一版就顺手验证一次真正上线的时候你只会看到一条顺畅的流水线而不是一夜之间冒出来的一堆报错。把这件“小事”放在 Phase A 的第三步恰恰是后面所有部署步骤能顺利推进的前提。
