年初接了个部署需求一个以Swin Transformer结构为主干的视觉模型训练用PyTorch最终要落到一台只有CPU的生产服务器上。模型参数量不算夸张但直接拿PyTorch推理单张图稳定跑在200ms上下离业务要求的100ms以内差着一截。圈里做部署的应该都清楚能走的路就那么几条ONNX Runtime、OpenVINO、TensorRT、NCNN而放到最前面试的通常是ONNX——生态最通用模型转换一次后续切不同推理后端都方便。真正动手才发现把一条简单的ResNet转成ONNX可能十分钟搞定但换成Swin Transformer这种带复杂窗口划分、相对位置编码、动态shape的模型报错能一路从torch导出阶段蔓延到ONNX Runtime加载阶段。这篇文章就把我这次复杂网络转ONNX的完整过程写下来包括一个个报错的出现场景、排查链路和最终修复代码以及转完之后真实的推理加速数据。反正我自己再回头看当初要是有这么一篇记录至少能少熬两个通宵。1. 为什么我坚持把PyTorch模型转成ONNX动机与现实收益1.1 场景尴尬PyTorch推理撑不起生产环境先说清楚我遇到的实际问题。模型在GPU上训练精度没问题但生产服务器没有GPU只能用CPU扛。PyTorch在CPU上的推理效率说实话一直不太行。原因在于PyTorch的Eager模式是一步一步调度算子每个算子都有Python层的开销小算子在CPU上频繁切换反而更慢虽然torch.compile在2.0之后解决了一部分问题但对于Swin Transformer这种结构复杂、包含大量切分重排操作的网络torch.compile的编译时间和兼容性又成了新问题。而我需要的是一个不依赖特定硬件、跨平台、能长期稳定运行的部署方案。ONNX是开放格式转出来之后可以在ONNX Runtime上跑之后如果业务量上来同一份ONNX还可以丢给TensorRT或OpenVINO继续优化属于性价比很高的选择。1.2 ONNX加速的本质图优化、算子融合与轻量运行时很多人以为ONNX加速是玄学或者以为转换本身就能让模型变快。其实ONNX格式本身不加速加速的是ONNX Runtime这个推理引擎。ONNX Runtime拿到模型后会先做一次完整的图优化把连续的ConvBNRelu融合成一个FusedConv把可以合并的Reshape/Transpose链合并掉把常量折叠提前算好。底层还针对CPU做了指令集优化。更关键的是整个推理过程没有Python层的反复调度每一个算子都是C层面的执行省掉的那部分开销在小算子密集的网络里非常可观。我用一个小实验验证过同一个ResNet50输入224x224图像PyTorch CPU单线程推理约66msONNX Runtime CPU单线程约42ms提速约1.57倍。换成结构更复杂、小算子更多的Swin-like网络提速会更明显。这也是为什么我坚持要转ONNX——不是跟风是确确实实能解决生产环境的问题。1.3 什么样的模型适合转ONNX什么样不建议这里先说句实话不是所有模型都适合转ONNX。含大量动态控制流if依赖tensor值、while循环次数依赖输入的网络转起来很痛苦。用了官方ONNX导出暂不支持的自定义算子比如某些3D网格采样、特殊池化需要自己注册符号。训练和部署完全在GPU上且GPU是N卡TensorRT提速通常比ONNX Runtime GPU更狠。模型结构极其简单纯MLP转换意义不大PyTorch本身也不慢。反过来说如果你要部署的目标环境是CPU服务器、边缘设备、跨平台场景或者要在多个推理框架之间切换那转ONNX基本是必经之路。这篇文章后面所有排错经验都是基于一个前提模型结构复杂但业务上非转不可。2. 导出前的准备版本搭配、导出参数与动态轴2.1 环境版本怎么搭才不容易踩坑这次踩坑给我最大的教训之一就是先确认版本对应关系。PyTorch、onnx、onnxruntime三个库的版本是互相影响的新旧混搭特别容易出一些莫名其妙的问题。我最终稳定使用的版本组合如下供参考组件本次使用版本说明PyTorch2.0.12.x系列均可1.13也能跑通大部分导出onnx1.14.0低于1.13有部分算子定义不全onnxruntime1.15.1CPU版本依赖较少torchvision0.15.2如果模型加载预训练权重需要安装方式直接pip就能搞定不需要额外编译。建议在单独的虚拟环境里做转换避免和训练环境里各种库相互污染。提示如果在conda环境里同时装了cuda版PyTorch和CPU版ONNX Runtime注意不要混着跑同一个进程。转换和推理最好脚本分离。2.2 torch.onnx.export的核心参数逐个说清导出一个模型核心函数是torch.onnx.export。用对了参数能省掉一半问题。我通常这样写import torch def export_to_onnx(model, dummy_input, export_pathmodel.onnx): model.eval() torch.onnx.export( model, (dummy_input,), export_path, export_paramsTrue, opset_version16, do_constant_foldingTrue, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, ) print(fexported to {export_path})几个参数单独说一下opset_version尽量用16或以上。Swin Transformer里的roll、grid_sample等算子在opset 11下要么不支持要么行为有偏差。opset高不一定完美但支持面更广。do_constant_folding开启常量折叠能将部分固定shape的计算提前算好减小模型体积推理也略快。默认True一般不用动。input_names/output_names名字随意起但后面用ONNX Runtime推理时输入输出的key就是这里的名字所以提前定好。dynamic_axes指定哪些维度是动态的。最常见的是把batch维设成动态。2.3 dynamic_axes动态shape是加速最大的变数这一条是我后面所有shape相关报错的根源必须单独拎出来讲。如果你确定推理时batch永远是1最简单的方式是不设dynamic_axes让模型固定输入shape。固定shape意味着ONNX Runtime能做更激进的图优化很多算子无需动态推断推理速度通常更快报错也少。但业务场景往往要求batch可变或者输入图像尺寸大小不一。这时候就必须设置dynamic_axes。问题是一旦有了动态维度模型里的view、reshape、flatten、transpose这些操作就会在导出和运行时面临维度推断的挑战。Swin Transformer恰恰是这些操作的密集户窗口划分、窗口合并、相对位置索引生成每一步都是shape操作。所以我的建议是先固定batch1把模型导出跑通再开动态轴。分两步排查问题不会搅在一起。3. 复杂网络的报错实录五个典型案例的完整排查3.1 TracerWarning动态控制流留下的定时炸弹第一个遇到的报错严格来说不是红字异常但比红字更阴险。导出时终端打出一长串警告TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect.我的模型里有一段类似这样的代码def forward(self, x): if x.shape[-1] 16: x self.forward_large(x) else: x self.forward_small(x) return xPyTorch导出采用追踪模式tracing也就是用dummy_input真实跑一遍模型记录所有算子调用。问题在于如果你的代码里有if条件依赖tensor的shape或数值trace时只会记录条件当前成立的那条分支另一条分支永远不会被记录。模型导出时看着正常但一旦推理时输入尺寸变化走了另一条分支行为就和预期完全不一致。这个问题在Swin Transformer里很隐蔽。它的窗口划分函数会根据输入尺寸决定是否padding到窗口大小的整数倍def forward(self, x, H, W): pad_l pad_t 0 pad_r (self.window_size[1] - W % self.window_size[1]) % self.window_size[1] pad_b (self.window_size[0] - H % self.window_size[0]) % self.window_size[0] if pad_r 0 or pad_b 0: x F.pad(x, (0, pad_r, 0, pad_b)) ...pad_r和pad_b在trace时如果是0if不成立padding分支就不会被记录。但实际部署时输入尺寸一变padding分支应该执行却没有执行结果就是张量shape直接对不上推理结果完全错误。修复方式把动态分支尽量改成无分支写法例如用torch.where替代if-else。padding这种可以用固定的pad值或者用F.pad的常量参数配合计算逻辑让导出时路径固定。如果控制流实在无法消除考虑用torch.jit.script脚本化脚本化能保留控制流结构但代价是有些Python操作不支持。我这个模型最后是把动态padding逻辑挪到了数据预处理阶段进模型的图像统一resize到窗口倍数尺寸模型内部的行分支就消除了。这是治本的办法。3.2 算子导出失败grid_sample和窗口划分第二个报错是红字异常发生在导出过程中RuntimeError: Exporting the operator aten::grid_sampler to ONNX opset version 11 is not supported. Please feel free to request support or submit a PR.原因很好理解模型里用torch.nn.functional.grid_sample做了特征采样而这个算子只在较新的opset里才支持。opset 11不支持导出直接不认。这个报错的排查思路很简单但也最容易疏忽先看报错里提示的算子名aten::grid_sampler定位到代码里对应的grid_sample调用。然后看当前导出用的opset版本如果低于16先尝试升级opset。如果升级opset后依然不支持才考虑自定义符号或者改写实现。我当时的做法torch.onnx.export( model, (dummy_input,), model.onnx, opset_version16, ... )在opset 16下直接通过。grid_sample这种算子其实ONNX官方是支持的只是需要opset高一点。类似的还有torch.rollSwin Transformer相对位置偏移里会用到在低opset下也会报aten::roll不支持。遇到这类问题第一反应都应该是先提opset版本而不是急着改算子。3.3 非常量属性动态pads为什么导不出去即将重提另一个经典报错RuntimeError: Failed to export an ONNX attribute pads, since its not constant, please try to make things traceable.这个报错出现在我把模型改成固定输入尺寸之后但内部有些padding计算仍然是动态的。比如pad_h some_dynamic_value # 依赖前面tensor的shape x F.pad(x, (0, pad_h, 0, pad_w))F.pad在ONNX中对应的算子Pad有一个pads属性要求是编译期常量。如果pads是运行时才知道的值PyTorch导出时就无法把它编码进ONNX图的attribute里直接报错。这种报错的核心原因就一句话ONNX的attribute必须静态动态的只能是tensor输入不能是attribute。凡是在Python层直接作为参数传给算子的值trace时必须推断为具体数值。解决方式有几种我最终采用的方法是使用torch.nn.functional.pad时把pad值写成const常量或者把动态pad逻辑挪到网络外面去。如果确实需要在网络内部处理可以考虑用一个支持动态padding的自定义算子但那会增加部署复杂度不值得。另外如果padding量确实依赖上一层输出shape可以尝试用torch.reshape配合torch.cat去模拟padding这样不是修改attribute而是新建tensorONNX只当作普通的concat操作处理不涉及常量属性问题。3.4 Reshape维度推断错误动态shape下最容易炸的地方这个报错简直是复杂网络转ONNX的噩梦导出时RuntimeError: The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 0或者导出成功了但在ONNX Runtime加载推理时报[ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Reshape node. Name:Reshape_126 Status Message: /onnxruntime/.../reshape_helper.h:... org.onnxruntime ORT FAIL : Contains negative values but also multiple 0 values.这类报错的根子在模型的view/reshape操作上。Swin Transformer做窗口划分时典型代码是这样的B, H, W, C x.shape x x.view(B, H // window_size, window_size, W // window_size, window_size, C)当H、W是动态维度时view需要PyTorch在导出时确定每个维度的具体值。如果某些维度是动态的torch.onnx.export试图把view转换成一个Reshape节点但Reshape的target shape里只要混入了-1和0ONNX Runtime的shape推断就会出问题——它不知道-1应该推断成多少或者shape里有0但实际输入对应维度不是0。我的排查过程先用onnx.checker.check_model检查导出后的模型是否通过校验。通过只能说明结构合法不代表运行不报错。然后再用onnxruntime跑一次推理看报错的具体node名。找到model.onnx里对应这个node的位置用Netron可视化关联回源码里的view调用。最终修复方式把view改成reshape。reshape在ONNX里同样对应Reshape node但它对shape的推断更宽松不是所有场景都有效。更稳妥的是避免在动态shape下使用view 手动指定维度改用torch.reshape并传入从x.shape动态获得的维度值B, H, W, C x.shape x torch.reshape(x, (B, H // wh, wh, W // ww, ww, C))最彻底的办法是让H、W变成固定值我在3.1里把图像预处理改成了固定尺寸后面这种问题就不再出现。实测下来固定H、W后ONNX Runtime的图优化效果最好Reshape报错几乎绝迹。3.5 ONNX Runtime加载验证阶段的新问题模型导出成功、检查通过但ONNX Runtime一加载就报错这种情况也很常见。我当时遇到的是[ONNXRuntimeError] : 1 : FAIL : Could not find an implementation for the node GatherElements_...问题定位模型里有若干torch.gather操作导出成GatherElements算子而当前ONNX Runtime版本对该算子的CPU实现有shape限制。排查思路是逐段注释代码缩小范围。最后我采用两种方案之一升级onnxruntime版本到1.15大部分算子实现都补全了。如果升级后仍不支持就在导出模型时把gather改写成等价的mulsum或one_hotmatmul组合。这个属于绕道实现能不出就别出。此外还有个经常遇到的是输出输出类型不匹配问题RuntimeError: Only tensors, lists, tuples, and dicts are supported as inputs/outputs of traced functions. Got unsupported type.比如模型forward返回了一个namedtuple或自定义类对象ONNX不认识。解决方式很简单把它改成tuple或dict返回。4. 验证与加速实测精度比对和真实的提速数据4.1 输出一致性校验转出来的ONNX不能只保证能导出、能加载还得保证输出结果和PyTorch一致。我写了一段标准的比对代码import numpy as np import onnxruntime as ort import torch def compare_pytorch_onnx(model, onnx_path, dummy_input, atol1e-4, rtol1e-4): model.eval() with torch.no_grad(): torch_outputs model(dummy_input) if isinstance(torch_outputs, (tuple, list)): torch_outputs [t.numpy() for t in torch_outputs] else: torch_outputs [torch_outputs.numpy()] so ort.SessionOptions() sess ort.InferenceSession(onnx_path, sess_optionsso, providers[CPUExecutionProvider]) ort_inputs {sess.get_inputs()[0].name: dummy_input.numpy()} ort_outputs sess.run(None, ort_inputs) for i, (torch_out, ort_out) in enumerate(zip(torch_outputs, ort_outputs)): np.testing.assert_allclose(torch_out, ort_out, rtolrtol, atolatol) print(foutput[{i}] shape{torch_out.shape}, max diff{np.abs(torch_out - ort_out).max():.6f})需要注意比对时用同一份dummy input且模型要切成eval()模式避免dropout/batchnorm差异。我的模型在atol1e-4下所有输出最大误差在1e-5量级说明转换精度损失可以忽略。4.2 CPU推理加速实测ONNX Runtime到底快多少验证通过之后我做了完整的推理延迟对比统一在CPU服务器上跑单线程、单图推理方式平均延迟相对PyTorch加速比PyTorch CPU (Eager)196ms1.0xONNX Runtime CPU (固定shape)84ms2.33xONNX Runtime CPU (动态batch)102ms1.92xONNX Runtime CPU int8动态量化63ms3.11x看到这个数据说实话有点惊讶。固定shape比动态shape快大约21%int8量化又在固定shape基础上再快25%。对生产环境来说这个收益非常香。但这里要明确一点加速效果因网络而异。Swin Transformer这种算子密集、内存重排多的网络ONNX Runtime的图优化收益很大如果是纯卷积堆叠的网络加速比可能没那么夸张但通常也能到1.5倍以上。4.3 动态shape与固定批次性能取舍既然对比数据摆出来了动态shape慢得多那到底要不要动态我的建议是服务端推理如果batch1直接固定shape别犹豫。需要batch可变可以只把batch设动态H、W保持固定。图像尺寸可变如果能统一resize到固定尺寸优先固定如果业务上无法接受才把H、W也动态化但要接受性能和稳定性的损耗。动态维度的数量直接影响ONNX Runtime的优化空间。每多一个动态维度一些算子就无法做静态shape推断内存分配和kernel选择都会变保守。反正能固定的维度尽量固定。5. 再往前一步int8量化和进一步加速的尝试5.1 dynamic quantization的踩坑记录当CPU推理从196ms优化到84ms之后我又动了量化的心思。onnxruntime的量化分为动态量化和静态量化。动态量化实现最简单不需要校准数据几行代码就能跑from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model.onnx, model_int8.onnx, weight_typeQuantType.QUInt8, )这个过程中我也踩了坑。第一次跑完量化模型加载没问题但运行时报[ONNXRuntimeError] : 1 : FAIL : Could not find an implementation for the node MatMulInteger...原因是量化后有些整数内核在CPU上不支持某个特定shape的MatMul整形内核。排查后发现是模型里某个动态shape的MatMul在量化后没有对应的kernel实现。解决方式把对应MatMul的量化和反量化配置手动关闭nodes_to_exclude参数指定跳过某些节点。或者把量化后仍然报错的输入shape固定下来再导出一次。如果只是部分算子不支持可以尝试不同量化类型如QuantType.QInt8和QUInt8交叉验证。最后我在quantize_dynamic中指定了排除某些节点成功跑通精度也从原先的1e-5误差放大到了约1e-2对视觉检测任务来说完全可用但严格的说还是要用任务指标比如mAP、准确率去评估不能只看输出误差。5.2 加速的边界与部署权衡做完int8量化后单张图63ms已经能满足业务需求。但如果还想继续压还有几条路静态量化需要采集校准数据集精度保持通常比动态量化好但操作复杂度上一个台阶。换推理后端同一份ONNX可以转OpenVINOCPU上有额外收益如果将来有GPUTensorRT还能再快。多线程优化设置SessionOptions.intra_op_num_threads和inter_op_num_threads但要注意线程数和延迟的关系不是越多越好。我自己测过4线程下ONNX Runtime的延迟从84ms降到了32ms但CPU占用率大幅上升服务端并发场景下需要考虑整机吞吐不能只看单次延迟。加速到最后会有边际效应从196ms到84ms是重构级的提升从84ms到63ms是量化收益再往后花大精力换来5-10ms的提升性价比就低了。这一步我选择了收手把重心放到部署稳定性和监控上。写在最后这次复杂网络从PyTorch转ONNX的经历让我最深的感受是排错的过程比最终结果更能让人成长。遇到报错先不要急着搜答案养成从报错信息反推根因的习惯——先看哪个算子、哪个维度、哪段代码再对照ONNX导出机制判断是属性问题、算子支持问题还是动态shape问题很多坑都能自己解。我踩过的坑里面动态控制流是最隐蔽的非常量属性是最容易理解的动态shape的Reshape是耗时最长的opset版本则是第一个该检查的。这几类问题在任何一个复杂网络转换里都有可能遇到。希望这篇记录能帮你少走一点弯路。如果你也在转一个结构更复杂的网络可以从opset和固定shape这两个地方开始排查八成能在前半小时内找到突破口。