本系列第一篇优雅地解决 onnx 算子不兼容 —— 算子注册讲的是导出前PyTorch → ONNX 阶段如何通过注册 symbolic 函数解决算子不兼容的问题。本系列第二篇优雅地改造 onnx 计算图 —— 算子替换与子图融合上讲的是导出后ONNX 文件已经生成了如何用onnx_graphsurgeon工具修改计算图完成算子替换与子图融合是解决算子不高效不兼容问题的基本技能与重要手段。本篇是本系列第三篇针对前两篇文章中的理论难点实践难点和痛点进行详细剖析。目录一、怎么知道张量的名字定位子图的四种方法1.1 方法一Netron 可视化最推荐90% 的情况用这个1.2 方法二用 gs 打印所有张量1.3 方法三遍历节点打印1.4 方法四按算子类型过滤大模型常用二、等价性验证替换前后结果必须一致2.1 为什么必须验证2.2 验证的三个层次2.3 验证代码模板2.4 数值精度问题严格等价 vs 近似等价三、深度拓展那些你可能困惑的问题3.1 ONNX 文件本身是壳子还是有底层实现3.2 标准算子 vs 自定义算子3.3 算子的数学语义在哪里定义3.3 两种注册的区别3.4 恒等映射存在的意义是什么3.5 ONNX 层面融合 vs TensorRT 层面融合会不会做无用功3.5.1 分层优化各管一段3.5.2 TensorRT 会自动融合哪些3.5.3 哪些必须在 ONNX 层面做3.6 未知的、带参数的算子不好替换怎么办3.6.1 先分清是哪种情况3.6.2 六种应对策略从优到劣策略一导出时就避免最优先成本最低策略二算子注册3.4 的方法最优雅策略三onnx_graphsurgeon 自动替换写模式匹配脚本策略四用 onnxsim 等工具自动简化策略五TensorRT Plugin最后手段最灵活策略六重新训练万不得已成本最高总结一、怎么知道张量的名字定位子图的四种方法前面两个案例里我们都是直接写tensors[identity_output_0]这样的代码。但你可能会问我怎么知道这些张量叫什么名字这是初学者最常问的问题。实际工作中你拿到的是别人导出的 ONNX 模型不可能知道张量名字。必须先探查再写替换代码。这里介绍四种探查方法前两种最常用。1.1 方法一Netron 可视化最推荐90% 的情况用这个Netron 是专门查看神经网络模型结构的可视化工具支持 ONNX、TensorRT、TensorFlow、PyTorch 等几乎所有格式。安装和使用pipinstallnetron netron model.onnx运行后浏览器会自动打开你会看到完整的计算图每个方框是一个算子节点上面写着 op 类型连线是张量点击连线可以看到张量的名字、形状、数据类型点击节点可以看到它的输入张量名字、输出张量名字、属性值或者可以不用安装直接在浏览器上打开https://netron.app/并将onnx文件拖入即可。定位子图的方法在图里找到你要替换的子图比如 LayerNorm 的那一堆算子点击子图最前面的那个节点看它的第一个输入张量叫什么名字——这就是子图的输入点击子图最后面的那个节点看它的输出张量叫什么名字——这就是子图的输出记下这些名字写到代码里这是最直观、最不容易出错的方法工作中 90% 的情况用 Netron 就够了。1.2 方法二用 gs 打印所有张量如果不想装 Netron或者需要在脚本里自动探查可以用 gs 打印所有张量importonnximportonnx_graphsurgeonasgs graphgs.import_onnx(onnx.load(model.onnx))tensorsgraph.tensors()print( 所有张量 )forname,tensorintensors.items():tensor_typetype(tensor).__name__# Variable 或 Constantprint(f{name:40s}|{tensor_type:8s}| dtype{tensor.dtype}| shape{tensor.shape})输出如下这样你能看到所有张量的名字、类型、形状从中挑出你需要的。1.3 方法三遍历节点打印如果想按节点来看而不是按张量用这段print( 所有节点 )fori,nodeinenumerate(graph.nodes):input_names[inp.nameforinpinnode.inputs]output_names[out.nameforoutinnode.outputs]print(f节点{i}: op{node.op:15s}| 输入{input_names}| 输出{output_names})输出如下这样能清楚看到每个节点的上下游连接关系定位子图边界非常方便。1.4 方法四按算子类型过滤大模型常用如果模型很大比如 Swin Transformer 有几百个节点全部打印出来看不过来可以按 op 类型过滤# 只找 ReduceMean 和 Div 节点LayerNorm 子图里的关键节点fornodeingraph.nodes:ifnode.opin(ReduceMean,Div,Mul,Add):print(fop{node.op})print(f 输入:{[inp.nameforinpinnode.inputs]})print(f 输出:{[out.nameforoutinnode.outputs]})输出如下实际工作中的标准流程先用 Netron 整体看一下模型结构找到要替换的子图大概在什么位置然后用 gs 打印相关节点和张量确认准确的名字最后写替换代码。二、等价性验证替换前后结果必须一致算子替换不是换完就可以。神经网络是一个整体局部替换如果不等价整个网络直接作废。所以替换后必须验证等价性。2.1 为什么必须验证即使你认为替换前后是数学等价的也可能因为各种原因出问题张量名字写错了接错了输入输出属性值传错了比如 epsilon、axisscale 和 bias 的值提取错了算子的数学语义理解有误比如某个算子的边界处理和你想的不一样浮点计算顺序不同导致的数值差异所以验证是必须的不能靠我觉得等价就跳过。2.2 验证的三个层次层次验证什么怎么验证什么时候做算子级替换的子图本身是否等价单独给子图喂输入比较替换前后的输出开发调试时网络级整个模型的输出是否一致用同一批输入比较整个模型的最终输出每次替换后必做业务级最终任务指标是否下降用真实测试集跑看 mAP/准确率/精度等上线前必做工作中至少要做到网络级验证业务级验证最好也做尤其是精度敏感的任务。2.3 验证代码模板这是我实际工作中常用的验证代码可以直接用importnumpyasnpimportonnxruntimedefcompare_models(model_a_path,model_b_path,test_input,input_nameinput0,atol1e-5,rtol1e-5): 比较两个 ONNX 模型在相同输入下的输出是否一致 Args: model_a_path: 原始模型路径 model_b_path: 修改后模型路径 test_input: 测试输入numpy 数组 input_name: 输入张量名字 atol: 绝对误差容忍度 rtol: 相对误差容忍度 Returns: bool: 是否等价 sess_aonnxruntime.InferenceSession(model_a_path)sess_bonnxruntime.InferenceSession(model_b_path)output_asess_a.run(None,{input_name:test_input})[0]output_bsess_b.run(None,{input_name:test_input})[0]# 逐元素比较|a - b| atol rtol * |b|is_closenp.allclose(output_a,output_b,atolatol,rtolrtol)max_diffnp.max(np.abs(output_a-output_b))mean_diffnp.mean(np.abs(output_a-output_b))print(f形状: A{output_a.shape}, B{output_b.shape})print(f是否等价:{is_close})print(f最大误差:{max_diff:.8f})print(f平均误差:{mean_diff:.8f})ifnotis_close:# 打印误差最大的位置方便调试diffnp.abs(output_a-output_b)idxnp.unravel_index(np.argmax(diff),diff.shape)print(f误差最大位置:{idx}, A{output_a[idx]}, B{output_b[idx]})returnis_close调用代码的实现if__name____main__:# 测试输入test_inputnp.random.rand(1,3,5,5).astype(np.float32)# 比较模型is_closecompare_models(model_a_path../models/sample-ln-before.onnx,model_b_path../models/sample-ln-after.onnx,test_inputtest_input)运行输出2.4 数值精度问题严格等价 vs 近似等价不是所有替换都能做到输出完全一样要分情况情况例子误差范围能不能接受严格数学等价MinMax → Clip、Identity 删除0完全一样✅ 放心用近似等价浮点误差不同计算顺序导致的误差比如 (ab)c vs a(bc)1e-7 ~ 1e-5✅ 一般可以接受功能等价但有精度损失量化、把复杂算子替换成简化版本1e-3 ~ 1e-1 或更大⚠️ 需要评估对任务的影响前两种可以放心替换第三种需要谨慎评估。工作中大部分算子融合ConvBN、LayerNorm 融合都是严格等价或近似等价的。误差容忍度的选择一般任务atol1e-5, rtol1e-5精度敏感任务如人脸识别、医疗影像atol1e-6, rtol1e-6量化后模型atol1e-2, rtol1e-2量化本身就有精度损失三、深度拓展那些你可能困惑的问题这部分整理了学习过程中最容易困惑的几个问题做深度拓展。3.1 ONNX 文件本身是壳子还是有底层实现答案ONNX 文件本身就是一个壳子不包含任何算子的底层算法实现。ONNX 文件是 protobuf 格式的二进制文件里面只存两类东西计算图的结构“图纸”节点列表每个节点的 op_type、输入张量名字、输出张量名字、属性张量信息名字、形状、数据类型连接关系哪个节点的输出连到哪个节点的输入权重数据“家具”模型训练好的参数以常量张量的形式存在文件里比如 Conv 的 weight、biasGemm 的 weight、biasLayerNorm 的 scale、bias这些是具体的数值不是代码ONNX 文件里绝对不存❌ CPU 上怎么用 SIMD 指令实现 Conv❌ GPU 上怎么写 CUDA kernel 实现 Relu❌ Clip 的逐元素截断逻辑怎么用代码写❌ 任何函数、类、可执行代码具体实现由推理后端提供onnxruntime有自己的 CPU/GPU 实现TensorRT有自己的 CUDA kernel 实现ncnn、MNN、OpenVINO各自有自己的实现当你用 onnxruntime 推理时onnxruntime 读到图里有个 Clip 节点就去自己的算子库里找到 Clip 的实现然后执行。这个实现过程和 ONNX 文件本身完全无关。3.2 标准算子 vs 自定义算子标准算子Clip、Conv、Relu自定义算子custom::xxxONNX 声明有有推理后端实现内置了onnxruntime/TensorRT 都有没有需要自己写 Plugin声明后能直接跑吗能不能只是空壳所以你在 gs 里创建一个 Clip 节点 → onnxruntime/TensorRT 读到后去自己的算子库里找到 Clip 的实现 → 能跑你在 gs 里创建一个custom::my_op节点 → onnxruntime/TensorRT 读到后发现自己的算子库里没有这个算子 → 报错这就是为什么后面学 TensorRT 部署的时候自定义算子还需要写 Plugin——ONNX 层面只是声明了节点TensorRT 层面还没有实现需要你自己补 CUDA kernel。3.3 算子的数学语义在哪里定义你可能会问既然 ONNX 不存实现那Conv 怎么算、LayerNorm 的公式是什么这些是在哪里规定的总不能每个后端自己随便实现吧**答案在 ONNX 官方的Operators.md文档里。**D这个文档是 ONNX 算子的圣经每个算子都有严格的数学定义版本历史since version简要描述Summary数学公式用文字或公式描述属性Attributes名字、类型、描述、默认值输入Inputs名字、类型、是否可选输出Outputs名字、类型类型约束Type Constraints示例Examples为什么需要这么严格的定义你可能觉得卷积、池化这些学过深度学习的都知道还用得着官方定义吗恰恰是因为看起来都知道才需要严格定义。不同框架对同一个算子的细节处理可能不一样如果没有统一规范跨框架部署就会出问题。举几个典型的细节坑算子容易有歧义的细节不同框架的差异Conv/MaxPoolpadding 的边界处理、ceil_mode输出形状向上取整还是向下取整PyTorch 的 ceil_modeTrue/False、TensorFlow 的 SAME/VALIDBatchNormepsilon 加在方差前还是方差后训练和推理的区别epsilon 默认值不同PyTorch 1e-5TensorFlow 1e-3Resize插值算法的具体定义坐标变换公式align_cornersPyTorch 的 align_cornersTrue/False 结果完全不同LayerNormaxis 是从前往后数还是从后往前数epsilon 加在哪里不同实现可能有细微差别Einsum等式的语法规则省略号怎么处理ONNX 有严格的等式解析规则NMS框的坐标格式iou 阈值的处理方式score 阈值不同框架的 NMS 实现细节差异很大如果没有官方的严格数学定义同一个 ONNX 模型在 onnxruntime 上跑和在 TensorRT 上跑结果可能不一样。有了规范所有后端都必须按照同一个数学定义去实现才能保证跨平台的一致性。数学语义 vs 实现算法还有一个关键区别官方说明的是数学语义结果应该是什么不是实现算法应该怎么算。比如 Conv数学语义输出的每个元素等于输入对应区域和权重的互相关求和 bias实现算法可以用 im2col GEMM、可以用 FFT、可以用 Winograd、可以用直接卷积……后端想怎么实现就怎么实现只要结果符合数学定义再比如 MatMul数学语义矩阵乘法实现算法CPU 上用 OpenBLAS/MKLGPU 上用 cuBLAS不同后端用不同的优化算法但结果一样ONNX 规范只约束结果对不对不约束算得快不快。这就是为什么不同推理后端的性能差异很大——它们用了不同的实现算法但都符合同一个数学语义规范。3.3 两种注册的区别很多人刚学的时候会把本系列第一篇的算子注册和第二篇的gs.Graph.register()“搞混因为它们都用了装饰器都叫注册”。但它们是完全不同阶段、完全不同目的的两件事。装饰器_onnx_symbolic/register_custom_op_symbolicgs.Graph.register()本质注册PyTorch 算子 → ONNX 算子的映射规则注册gs 工具的便捷方法语法糖阶段PyTorch → ONNX导出时ONNX 图已存在用 gs修改时注册给谁PyTorch 的 ONNX 导出器symbolic_registrygs.Graph这个 Python 类注册的内容PyTorch 算子到 ONNX 算子的翻译规则创建节点的便捷辅助函数为什么需要导出器不认识这个 PyTorch 算子不知道怎么转纯粹为了写代码方便不注册也能用 gs.Node 创建不注册会怎样导出直接报错无法识别的算子代码写起来啰嗦但功能完全不受影响和 ONNX 算子支持的关系间接相关ONNX 支持但导出器不知道怎么映射完全无关ONNX 本来就支持gs 本来就能创建用比喻来理解上篇文章的注册你家里来了个亲戚PyTorch 算子门卫导出器不认识他不让进导出报错。你去门卫那里登记一下注册映射规则说这个亲戚对应 ONNX 王国里的 Asinh 公民门卫就放行了。这篇文章的注册你已经在 ONNX 王国里了要盖房子创建节点。工具有两种用法一种是每次都手动搬砖gs.Node(...)另一种是你给工具箱装个快捷按钮gs.Graph.register()一键就能放好 Min 砖块。装不装快捷按钮房子都能盖只是装了更方便。一句话总结本篇的注册不注册也能用只是啰嗦。3.4 恒等映射存在的意义是什么Identity 算子的输出等于输入数学上什么都不做。那它为什么存在Identity 在工程上有几个非常实际的用途用途说明子图边界标记在复杂图里Identity 可以清晰地标出一个子图的输入输出边界方便替换/提取子图张量重命名ONNX 里张量名字很重要很多工具按名字引用。Identity 可以给张量改名字——输入叫 A输出叫 BB 就是 A 的别名调试探针在 Netron 里看复杂模型时插一个 Identity 可以让某个中间张量更醒目方便定位和观察框架导出的产物PyTorch/TensorFlow 导出 ONNX 时经常自动插入 Identity比如模型里写了 xx、或者 trace 过程中产生的冗余节点图结构调整某些推理后端对图的输入输出格式有要求Identity 可以用来调整连接关系回到gs_replace_op.py这个案例图结构是这两个 Identity 是故意放的边界标记Identity0 的输出 MinMax 子图的输入边界Identity1 的输入 MinMax 子图的输出边界这样替换的时候边界非常清楚子图就是两个 Identity 之间的 MinMax。替换成 Clip 后两个 Identity 保持不动输入输出的连接关系也很清楚。如果没有这两个 Identity图就是input0 → Min → Max → output子图边界就模糊了——Min 的输入就是整个图的输入Max 的输出就是整个图的输出。替换的时候需要直接操作图的整体输入输出教学上就不那么直观了。实际工作中怎么处理 Identity实际工作中Identity 大部分时候是冗余节点框架导出时自动产生的对计算没有任何贡献还会占一点点推理开销。通常直接删掉onnxsim工具会自动删掉 Identity 等冗余节点或者用 gs 手动清理但在子图替换、调试、张量重命名这些特定场景下Identity 是有用的工具——就像这个示例里一样它帮你清晰地框出要操作的范围。3.5 ONNX 层面融合 vs TensorRT 层面融合会不会做无用功这是一个非常实际的工程问题我在 ONNX 层面把 LayerNorm 的 10 个小算子融合成了 1 个但万一 TensorRT 在生成 engine 的时候也会自动融合这些小算子呢那我前置工作岂不是白做了答案不会。ONNX 层面融合和 TensorRT 层面融合是两个不同阶段的优化各管一段各有不可替代的价值。3.5.1 分层优化各管一段模型从 PyTorch 到最终跑起来要经过好几层优化PyTorch 模型 ↓ 【导出层】算子注册、opset 选择 ONNX 模型 ↓ 【ONNX 层】算子融合、子图替换、常量折叠、形状简化本文讲的 ONNX 模型优化后 ↓ 【TensorRT 构建层】算子融合、kernel 自动调优、内存优化、量化 TensorRT Engine ↓ 【运行时层】多流并行、CUDA Graph 实际推理每一层的优化能力和范围不同不是说后面一层做了前面一层就白做了。3.5.2 TensorRT 会自动融合哪些TensorRT 的优化器确实很强大会自动做很多融合融合模式TensorRT 会自动做吗Conv BN✅ 会构建时把 BN 参数合并到 Conv 的 weight/biasConv Bias ReLU/LeakyReLU/Sigmoid✅ 会融合成 CBR 单个 kernelMatMul Add 激活✅ 会全连接层融合简单的恒等变换Identity、Flatten✅ 可能被消除LayerNorm 的 10 算子组合❌大概率不会复杂的 attention 子图❌不会自定义的后处理逻辑❌不会TensorRT 的融合是基于固定模式匹配的——它知道Conv 后面跟 BN 可以融合但不知道这 10 个乱七八糟的算子组合起来等价于 LayerNorm。反向推断从一堆小算子反推出一个大算子在工程上非常难TensorRT 没有做这个。3.5.3 哪些必须在 ONNX 层面做情况为什么必须在 ONNX 层面做TensorRT 不支持的算子不替换就跑不通没得商量复杂归一化子图LN、IN、GNTensorRT 大概率不自动融合一个个 launch kernel 性能差复杂的 attention 子图TensorRT 认不出来需要替换成它能优化的结构动态形状改静态TensorRT 对动态形状支持有限需要在 ONNX 层面简化修改节点属性/输入比如改 Reshape 的 shape只能在 ONNX 层面做怎么判断需不需要在 ONNX 层面做不要上来就手动融合所有东西按这个流程判断先直接用 TensorRT 跑看能不能跑通如果报错不支持某个算子→ 必须在 ONNX 层面处理如果能跑通 → 进入第二步看性能profiling 找瓶颈trtexec--onnxmodel.onnx--saveEnginemodel.engine--dumpProfile如果发现某一组小算子各自占一点时间加起来总耗时很大 → 说明 TensorRT 没有融合需要在 ONNX 层面手动融合如果常见的 ConvBN 已经被融合成一个节点了 → 不用管看 TensorRT 的优化日志加--verbose看融合前后的层数变化就能大概知道 TensorRT 做了多少融合经验法则常见模式ConvBN、ConvReLU→ 不用TensorRT 自动做复杂子图LayerNorm、attention、后处理→ 要做TensorRT 大概率不自动融合不支持的算子 → 必须做即使 TensorRT 会做ONNX 层面做也有价值。有些情况下即使 TensorRT 能自动融合在 ONNX 层面先做一遍也不是完全没意义价值说明跨平台兼容性融合后的 ONNX 模型在 onnxruntime、ncnn、OpenVINO 等其他后端上也能受益不只是 TensorRT模型可读性融合后的模型用 Netron 打开更简洁调试、排查问题更方便加快 TensorRT 构建速度TensorRT 做融合也需要时间模式匹配、图遍历ONNX 层面先做好可以缩短引擎构建时间避免版本差异不同版本的 TensorRT 融合能力不同在 ONNX 层面手动融合可以保证在所有版本上都有稳定的好性能减少中间张量的内存占用融合后中间张量不需要物化到内存所以即使是 ConvBN 这种 TensorRT 会自动融合的用 onnxsim 先跑一遍也没坏处——onnxsim 是自动的不需要你手动写代码成本很低。3.6 未知的、带参数的算子不好替换怎么办前面的案例都是已知的、无参数的或参数简单的算子替换。但实际工作中你可能会遇到未知的、带可训练参数的算子手动不好替换。这时候怎么办3.6.1 先分清是哪种情况情况例子是不是问题标准算子带参数Conv、Linear、BatchNorm❌ 不是问题ONNX 本来就支持参数自动导出为常量张量自定义算子带可训练参数自定义 attention 模块、自定义卷积变体、自定义激活函数⚠️ 是问题需要额外处理标准算子组合太碎LayerNorm 导出成 10 多个算子✅ 可以自动融合不是大问题你问的应该是第二类——自定义的、带可训练参数的算子。3.6.2 六种应对策略从优到劣策略一导出时就避免最优先成本最低最好的问题是不存在的问题。在模型设计或导出阶段就避免自定义算子用 ONNX 支持的标准算子重新实现自定义算子比如自定义的 attention可以用MatMul Softmax MatMul组合实现比如自定义的激活函数看看能不能用标准算子组合出来大部分自定义算子都能拆成标准算子的组合选高版本 opset 导出opset 越高支持的算子越多比如 LayerNorm 在 opset 17 才成为标准算子导出时尽量选最新的稳定 opset避免用 PyTorch 的自定义 autograd Functiontorch.autograd.Function导出时最容易出问题如果必须用确保写了 symbolic 方法3.4 的算子注册这是成本最低的方案能在导出阶段解决就不要拖到后面。策略二算子注册3.4 的方法最优雅如果 PyTorch 能识别这个算子比如 torchvision 里的、或者你自己写了 torch.ops 注册的只是 ONNX 导出器不认识那就用 3.4 的方法注册 symbolic 函数。带参数也没关系——参数weight、bias作为输入张量传过去导出后自动变成 ONNX 里的常量张量。这是最优雅的解决方案不需要手动改 ONNX导出时自动处理。策略三onnx_graphsurgeon 自动替换写模式匹配脚本如果已经导出成 ONNX 了而且子图模式是固定的可以写脚本自动识别和替换不用手动定位# 遍历所有节点找到符合某种模式的子图自动替换fornodeingraph.nodes:ifnode.opReduceMean:# 检查后面是不是跟着 Sub → Pow → ReduceMean → Add → Sqrt → Div...# 如果是就替换成单个 LayerNormalization...带参数也没关系——参数是常量张量替换时直接把旧的常量张量接过来就行。如果新算子需要不同格式的参数用gs.Constant创建新的常量张量把旧参数转换后放进去。本文的案例是手动定位的因为教学目的实际工作中可以写自动化的模式匹配脚本。策略四用 onnxsim 等工具自动简化很多常见的子图融合不需要你手动写脚本现成的工具自动就做了onnxsimONNX Simplifier自动消除冗余节点、常量折叠、算子融合ConvBN 融合、ConvRelu 融合、Identity 删除、Shape 相关节点简化……一行命令python -m onnxsim input.onnx output.onnxonnxruntime 的优化推理时自动做图优化TensorRT 的优化构建引擎时自动做算子融合和优化工作中拿到一个 ONNX 模型先跑一遍 onnxsim80% 的冗余和碎算子就自动处理了。剩下 onnxsim 处理不了的再手动用 gs 替换。策略五TensorRT Plugin最后手段最灵活如果以上方法都不行比如自定义算子太特殊无法拆成标准算子组合也无法等价替换就在TensorRT 里写自定义 Plugin把整个自定义算子包括参数、前向计算逻辑封装成一个 TensorRT PluginONNX 里只需要声明这个自定义算子节点用 gs 或者导出时注册TensorRT 推理时遇到这个节点就调用你写的 Plugin 来计算这是最灵活但也是最麻烦的方法需要写 CUDA kernel门槛高。一般到这一步说明这个算子确实很特殊。策略六重新训练万不得已成本最高如果自定义算子的结构和标准算子差异太大无法等价替换也无法写 Plugin或者写 Plugin 成本太高最后才考虑重新训练用标准算子重新实现一个功能类似的模块加载预训练权重然后微调fine-tune让新模块的输出尽量接近旧模块这是成本最高的方案需要数据、训练资源、时间而且可能有精度损失。一般不到万不得已不用。是算子设计不合理吗是 PyTorch 不合理吗都不是。这是整个行业的普遍问题叫研究-部署鸿沟research-deployment gapPyTorch 的设计目标是灵活的研究框架——支持任意自定义算子、任意计算图这是它的核心优势。如果限制只能用 ONNX 支持的算子那研究就没法做了新想法怎么验证ONNX 的设计目标是通用的中间表示——它不可能覆盖所有可能的自定义算子只能定义一套标准算子集覆盖 95% 的常见场景。部署框架TensorRT 等的设计目标是高性能推理——它们只优化标准算子自定义算子需要用户自己写 Plugin。这不是谁的错是三个不同目标的框架之间的天然矛盾。好的做法是研究阶段灵活设计随便用自定义算子快速验证想法部署阶段做算子工程化——把自定义算子改写成标准算子组合、或者写 Plugin、或者微调产品意识强的团队在模型设计阶段就考虑部署友好性尽量用标准算子避免太特殊的自定义结构总结本文围绕 ONNX 算子替换与子图融合的实践难点展开。首先介绍了定位子图的四种方法Netron 可视化最直观、gs 打印张量适合脚本探查、遍历节点可看清连接关系、按算子类型过滤适合大模型并给出先 Netron 看结构、再 gs 确认名字、最后写代码的标准流程。其次强调替换后必须做等价性验证分算子级、网络级、业务级三个层次并提供了可直接复用的验证代码模板同时区分了严格等价、近似等价与精度损失三种情况。最后深度解答了常见困惑ONNX 只是壳子、算子语义由官方文档定义、标准算子与自定义算子的区别、Identity 的工程用途、ONNX 与 TensorRT 分层融合的价值以及面对未知带参算子时的六种应对策略帮助读者系统掌握算子替换的完整方法论。本系列第一篇优雅地解决 onnx 算子不兼容 —— 算子注册本系列第二篇优雅地改造 onnx 计算图 —— 算子替换与子图融合上如果觉得有帮助欢迎点赞❤️、收藏❤️、关注❤️有问题可以在评论区交流。
优雅地改造 onnx 计算图 —— 算子替换与子图融合（下）
本系列第一篇优雅地解决 onnx 算子不兼容 —— 算子注册讲的是导出前PyTorch → ONNX 阶段如何通过注册 symbolic 函数解决算子不兼容的问题。本系列第二篇优雅地改造 onnx 计算图 —— 算子替换与子图融合上讲的是导出后ONNX 文件已经生成了如何用onnx_graphsurgeon工具修改计算图完成算子替换与子图融合是解决算子不高效不兼容问题的基本技能与重要手段。本篇是本系列第三篇针对前两篇文章中的理论难点实践难点和痛点进行详细剖析。目录一、怎么知道张量的名字定位子图的四种方法1.1 方法一Netron 可视化最推荐90% 的情况用这个1.2 方法二用 gs 打印所有张量1.3 方法三遍历节点打印1.4 方法四按算子类型过滤大模型常用二、等价性验证替换前后结果必须一致2.1 为什么必须验证2.2 验证的三个层次2.3 验证代码模板2.4 数值精度问题严格等价 vs 近似等价三、深度拓展那些你可能困惑的问题3.1 ONNX 文件本身是壳子还是有底层实现3.2 标准算子 vs 自定义算子3.3 算子的数学语义在哪里定义3.3 两种注册的区别3.4 恒等映射存在的意义是什么3.5 ONNX 层面融合 vs TensorRT 层面融合会不会做无用功3.5.1 分层优化各管一段3.5.2 TensorRT 会自动融合哪些3.5.3 哪些必须在 ONNX 层面做3.6 未知的、带参数的算子不好替换怎么办3.6.1 先分清是哪种情况3.6.2 六种应对策略从优到劣策略一导出时就避免最优先成本最低策略二算子注册3.4 的方法最优雅策略三onnx_graphsurgeon 自动替换写模式匹配脚本策略四用 onnxsim 等工具自动简化策略五TensorRT Plugin最后手段最灵活策略六重新训练万不得已成本最高总结一、怎么知道张量的名字定位子图的四种方法前面两个案例里我们都是直接写tensors[identity_output_0]这样的代码。但你可能会问我怎么知道这些张量叫什么名字这是初学者最常问的问题。实际工作中你拿到的是别人导出的 ONNX 模型不可能知道张量名字。必须先探查再写替换代码。这里介绍四种探查方法前两种最常用。1.1 方法一Netron 可视化最推荐90% 的情况用这个Netron 是专门查看神经网络模型结构的可视化工具支持 ONNX、TensorRT、TensorFlow、PyTorch 等几乎所有格式。安装和使用pipinstallnetron netron model.onnx运行后浏览器会自动打开你会看到完整的计算图每个方框是一个算子节点上面写着 op 类型连线是张量点击连线可以看到张量的名字、形状、数据类型点击节点可以看到它的输入张量名字、输出张量名字、属性值或者可以不用安装直接在浏览器上打开https://netron.app/并将onnx文件拖入即可。定位子图的方法在图里找到你要替换的子图比如 LayerNorm 的那一堆算子点击子图最前面的那个节点看它的第一个输入张量叫什么名字——这就是子图的输入点击子图最后面的那个节点看它的输出张量叫什么名字——这就是子图的输出记下这些名字写到代码里这是最直观、最不容易出错的方法工作中 90% 的情况用 Netron 就够了。1.2 方法二用 gs 打印所有张量如果不想装 Netron或者需要在脚本里自动探查可以用 gs 打印所有张量importonnximportonnx_graphsurgeonasgs graphgs.import_onnx(onnx.load(model.onnx))tensorsgraph.tensors()print( 所有张量 )forname,tensorintensors.items():tensor_typetype(tensor).__name__# Variable 或 Constantprint(f{name:40s}|{tensor_type:8s}| dtype{tensor.dtype}| shape{tensor.shape})输出如下这样你能看到所有张量的名字、类型、形状从中挑出你需要的。1.3 方法三遍历节点打印如果想按节点来看而不是按张量用这段print( 所有节点 )fori,nodeinenumerate(graph.nodes):input_names[inp.nameforinpinnode.inputs]output_names[out.nameforoutinnode.outputs]print(f节点{i}: op{node.op:15s}| 输入{input_names}| 输出{output_names})输出如下这样能清楚看到每个节点的上下游连接关系定位子图边界非常方便。1.4 方法四按算子类型过滤大模型常用如果模型很大比如 Swin Transformer 有几百个节点全部打印出来看不过来可以按 op 类型过滤# 只找 ReduceMean 和 Div 节点LayerNorm 子图里的关键节点fornodeingraph.nodes:ifnode.opin(ReduceMean,Div,Mul,Add):print(fop{node.op})print(f 输入:{[inp.nameforinpinnode.inputs]})print(f 输出:{[out.nameforoutinnode.outputs]})输出如下实际工作中的标准流程先用 Netron 整体看一下模型结构找到要替换的子图大概在什么位置然后用 gs 打印相关节点和张量确认准确的名字最后写替换代码。二、等价性验证替换前后结果必须一致算子替换不是换完就可以。神经网络是一个整体局部替换如果不等价整个网络直接作废。所以替换后必须验证等价性。2.1 为什么必须验证即使你认为替换前后是数学等价的也可能因为各种原因出问题张量名字写错了接错了输入输出属性值传错了比如 epsilon、axisscale 和 bias 的值提取错了算子的数学语义理解有误比如某个算子的边界处理和你想的不一样浮点计算顺序不同导致的数值差异所以验证是必须的不能靠我觉得等价就跳过。2.2 验证的三个层次层次验证什么怎么验证什么时候做算子级替换的子图本身是否等价单独给子图喂输入比较替换前后的输出开发调试时网络级整个模型的输出是否一致用同一批输入比较整个模型的最终输出每次替换后必做业务级最终任务指标是否下降用真实测试集跑看 mAP/准确率/精度等上线前必做工作中至少要做到网络级验证业务级验证最好也做尤其是精度敏感的任务。2.3 验证代码模板这是我实际工作中常用的验证代码可以直接用importnumpyasnpimportonnxruntimedefcompare_models(model_a_path,model_b_path,test_input,input_nameinput0,atol1e-5,rtol1e-5): 比较两个 ONNX 模型在相同输入下的输出是否一致 Args: model_a_path: 原始模型路径 model_b_path: 修改后模型路径 test_input: 测试输入numpy 数组 input_name: 输入张量名字 atol: 绝对误差容忍度 rtol: 相对误差容忍度 Returns: bool: 是否等价 sess_aonnxruntime.InferenceSession(model_a_path)sess_bonnxruntime.InferenceSession(model_b_path)output_asess_a.run(None,{input_name:test_input})[0]output_bsess_b.run(None,{input_name:test_input})[0]# 逐元素比较|a - b| atol rtol * |b|is_closenp.allclose(output_a,output_b,atolatol,rtolrtol)max_diffnp.max(np.abs(output_a-output_b))mean_diffnp.mean(np.abs(output_a-output_b))print(f形状: A{output_a.shape}, B{output_b.shape})print(f是否等价:{is_close})print(f最大误差:{max_diff:.8f})print(f平均误差:{mean_diff:.8f})ifnotis_close:# 打印误差最大的位置方便调试diffnp.abs(output_a-output_b)idxnp.unravel_index(np.argmax(diff),diff.shape)print(f误差最大位置:{idx}, A{output_a[idx]}, B{output_b[idx]})returnis_close调用代码的实现if__name____main__:# 测试输入test_inputnp.random.rand(1,3,5,5).astype(np.float32)# 比较模型is_closecompare_models(model_a_path../models/sample-ln-before.onnx,model_b_path../models/sample-ln-after.onnx,test_inputtest_input)运行输出2.4 数值精度问题严格等价 vs 近似等价不是所有替换都能做到输出完全一样要分情况情况例子误差范围能不能接受严格数学等价MinMax → Clip、Identity 删除0完全一样✅ 放心用近似等价浮点误差不同计算顺序导致的误差比如 (ab)c vs a(bc)1e-7 ~ 1e-5✅ 一般可以接受功能等价但有精度损失量化、把复杂算子替换成简化版本1e-3 ~ 1e-1 或更大⚠️ 需要评估对任务的影响前两种可以放心替换第三种需要谨慎评估。工作中大部分算子融合ConvBN、LayerNorm 融合都是严格等价或近似等价的。误差容忍度的选择一般任务atol1e-5, rtol1e-5精度敏感任务如人脸识别、医疗影像atol1e-6, rtol1e-6量化后模型atol1e-2, rtol1e-2量化本身就有精度损失三、深度拓展那些你可能困惑的问题这部分整理了学习过程中最容易困惑的几个问题做深度拓展。3.1 ONNX 文件本身是壳子还是有底层实现答案ONNX 文件本身就是一个壳子不包含任何算子的底层算法实现。ONNX 文件是 protobuf 格式的二进制文件里面只存两类东西计算图的结构“图纸”节点列表每个节点的 op_type、输入张量名字、输出张量名字、属性张量信息名字、形状、数据类型连接关系哪个节点的输出连到哪个节点的输入权重数据“家具”模型训练好的参数以常量张量的形式存在文件里比如 Conv 的 weight、biasGemm 的 weight、biasLayerNorm 的 scale、bias这些是具体的数值不是代码ONNX 文件里绝对不存❌ CPU 上怎么用 SIMD 指令实现 Conv❌ GPU 上怎么写 CUDA kernel 实现 Relu❌ Clip 的逐元素截断逻辑怎么用代码写❌ 任何函数、类、可执行代码具体实现由推理后端提供onnxruntime有自己的 CPU/GPU 实现TensorRT有自己的 CUDA kernel 实现ncnn、MNN、OpenVINO各自有自己的实现当你用 onnxruntime 推理时onnxruntime 读到图里有个 Clip 节点就去自己的算子库里找到 Clip 的实现然后执行。这个实现过程和 ONNX 文件本身完全无关。3.2 标准算子 vs 自定义算子标准算子Clip、Conv、Relu自定义算子custom::xxxONNX 声明有有推理后端实现内置了onnxruntime/TensorRT 都有没有需要自己写 Plugin声明后能直接跑吗能不能只是空壳所以你在 gs 里创建一个 Clip 节点 → onnxruntime/TensorRT 读到后去自己的算子库里找到 Clip 的实现 → 能跑你在 gs 里创建一个custom::my_op节点 → onnxruntime/TensorRT 读到后发现自己的算子库里没有这个算子 → 报错这就是为什么后面学 TensorRT 部署的时候自定义算子还需要写 Plugin——ONNX 层面只是声明了节点TensorRT 层面还没有实现需要你自己补 CUDA kernel。3.3 算子的数学语义在哪里定义你可能会问既然 ONNX 不存实现那Conv 怎么算、LayerNorm 的公式是什么这些是在哪里规定的总不能每个后端自己随便实现吧**答案在 ONNX 官方的Operators.md文档里。**D这个文档是 ONNX 算子的圣经每个算子都有严格的数学定义版本历史since version简要描述Summary数学公式用文字或公式描述属性Attributes名字、类型、描述、默认值输入Inputs名字、类型、是否可选输出Outputs名字、类型类型约束Type Constraints示例Examples为什么需要这么严格的定义你可能觉得卷积、池化这些学过深度学习的都知道还用得着官方定义吗恰恰是因为看起来都知道才需要严格定义。不同框架对同一个算子的细节处理可能不一样如果没有统一规范跨框架部署就会出问题。举几个典型的细节坑算子容易有歧义的细节不同框架的差异Conv/MaxPoolpadding 的边界处理、ceil_mode输出形状向上取整还是向下取整PyTorch 的 ceil_modeTrue/False、TensorFlow 的 SAME/VALIDBatchNormepsilon 加在方差前还是方差后训练和推理的区别epsilon 默认值不同PyTorch 1e-5TensorFlow 1e-3Resize插值算法的具体定义坐标变换公式align_cornersPyTorch 的 align_cornersTrue/False 结果完全不同LayerNormaxis 是从前往后数还是从后往前数epsilon 加在哪里不同实现可能有细微差别Einsum等式的语法规则省略号怎么处理ONNX 有严格的等式解析规则NMS框的坐标格式iou 阈值的处理方式score 阈值不同框架的 NMS 实现细节差异很大如果没有官方的严格数学定义同一个 ONNX 模型在 onnxruntime 上跑和在 TensorRT 上跑结果可能不一样。有了规范所有后端都必须按照同一个数学定义去实现才能保证跨平台的一致性。数学语义 vs 实现算法还有一个关键区别官方说明的是数学语义结果应该是什么不是实现算法应该怎么算。比如 Conv数学语义输出的每个元素等于输入对应区域和权重的互相关求和 bias实现算法可以用 im2col GEMM、可以用 FFT、可以用 Winograd、可以用直接卷积……后端想怎么实现就怎么实现只要结果符合数学定义再比如 MatMul数学语义矩阵乘法实现算法CPU 上用 OpenBLAS/MKLGPU 上用 cuBLAS不同后端用不同的优化算法但结果一样ONNX 规范只约束结果对不对不约束算得快不快。这就是为什么不同推理后端的性能差异很大——它们用了不同的实现算法但都符合同一个数学语义规范。3.3 两种注册的区别很多人刚学的时候会把本系列第一篇的算子注册和第二篇的gs.Graph.register()“搞混因为它们都用了装饰器都叫注册”。但它们是完全不同阶段、完全不同目的的两件事。装饰器_onnx_symbolic/register_custom_op_symbolicgs.Graph.register()本质注册PyTorch 算子 → ONNX 算子的映射规则注册gs 工具的便捷方法语法糖阶段PyTorch → ONNX导出时ONNX 图已存在用 gs修改时注册给谁PyTorch 的 ONNX 导出器symbolic_registrygs.Graph这个 Python 类注册的内容PyTorch 算子到 ONNX 算子的翻译规则创建节点的便捷辅助函数为什么需要导出器不认识这个 PyTorch 算子不知道怎么转纯粹为了写代码方便不注册也能用 gs.Node 创建不注册会怎样导出直接报错无法识别的算子代码写起来啰嗦但功能完全不受影响和 ONNX 算子支持的关系间接相关ONNX 支持但导出器不知道怎么映射完全无关ONNX 本来就支持gs 本来就能创建用比喻来理解上篇文章的注册你家里来了个亲戚PyTorch 算子门卫导出器不认识他不让进导出报错。你去门卫那里登记一下注册映射规则说这个亲戚对应 ONNX 王国里的 Asinh 公民门卫就放行了。这篇文章的注册你已经在 ONNX 王国里了要盖房子创建节点。工具有两种用法一种是每次都手动搬砖gs.Node(...)另一种是你给工具箱装个快捷按钮gs.Graph.register()一键就能放好 Min 砖块。装不装快捷按钮房子都能盖只是装了更方便。一句话总结本篇的注册不注册也能用只是啰嗦。3.4 恒等映射存在的意义是什么Identity 算子的输出等于输入数学上什么都不做。那它为什么存在Identity 在工程上有几个非常实际的用途用途说明子图边界标记在复杂图里Identity 可以清晰地标出一个子图的输入输出边界方便替换/提取子图张量重命名ONNX 里张量名字很重要很多工具按名字引用。Identity 可以给张量改名字——输入叫 A输出叫 BB 就是 A 的别名调试探针在 Netron 里看复杂模型时插一个 Identity 可以让某个中间张量更醒目方便定位和观察框架导出的产物PyTorch/TensorFlow 导出 ONNX 时经常自动插入 Identity比如模型里写了 xx、或者 trace 过程中产生的冗余节点图结构调整某些推理后端对图的输入输出格式有要求Identity 可以用来调整连接关系回到gs_replace_op.py这个案例图结构是这两个 Identity 是故意放的边界标记Identity0 的输出 MinMax 子图的输入边界Identity1 的输入 MinMax 子图的输出边界这样替换的时候边界非常清楚子图就是两个 Identity 之间的 MinMax。替换成 Clip 后两个 Identity 保持不动输入输出的连接关系也很清楚。如果没有这两个 Identity图就是input0 → Min → Max → output子图边界就模糊了——Min 的输入就是整个图的输入Max 的输出就是整个图的输出。替换的时候需要直接操作图的整体输入输出教学上就不那么直观了。实际工作中怎么处理 Identity实际工作中Identity 大部分时候是冗余节点框架导出时自动产生的对计算没有任何贡献还会占一点点推理开销。通常直接删掉onnxsim工具会自动删掉 Identity 等冗余节点或者用 gs 手动清理但在子图替换、调试、张量重命名这些特定场景下Identity 是有用的工具——就像这个示例里一样它帮你清晰地框出要操作的范围。3.5 ONNX 层面融合 vs TensorRT 层面融合会不会做无用功这是一个非常实际的工程问题我在 ONNX 层面把 LayerNorm 的 10 个小算子融合成了 1 个但万一 TensorRT 在生成 engine 的时候也会自动融合这些小算子呢那我前置工作岂不是白做了答案不会。ONNX 层面融合和 TensorRT 层面融合是两个不同阶段的优化各管一段各有不可替代的价值。3.5.1 分层优化各管一段模型从 PyTorch 到最终跑起来要经过好几层优化PyTorch 模型 ↓ 【导出层】算子注册、opset 选择 ONNX 模型 ↓ 【ONNX 层】算子融合、子图替换、常量折叠、形状简化本文讲的 ONNX 模型优化后 ↓ 【TensorRT 构建层】算子融合、kernel 自动调优、内存优化、量化 TensorRT Engine ↓ 【运行时层】多流并行、CUDA Graph 实际推理每一层的优化能力和范围不同不是说后面一层做了前面一层就白做了。3.5.2 TensorRT 会自动融合哪些TensorRT 的优化器确实很强大会自动做很多融合融合模式TensorRT 会自动做吗Conv BN✅ 会构建时把 BN 参数合并到 Conv 的 weight/biasConv Bias ReLU/LeakyReLU/Sigmoid✅ 会融合成 CBR 单个 kernelMatMul Add 激活✅ 会全连接层融合简单的恒等变换Identity、Flatten✅ 可能被消除LayerNorm 的 10 算子组合❌大概率不会复杂的 attention 子图❌不会自定义的后处理逻辑❌不会TensorRT 的融合是基于固定模式匹配的——它知道Conv 后面跟 BN 可以融合但不知道这 10 个乱七八糟的算子组合起来等价于 LayerNorm。反向推断从一堆小算子反推出一个大算子在工程上非常难TensorRT 没有做这个。3.5.3 哪些必须在 ONNX 层面做情况为什么必须在 ONNX 层面做TensorRT 不支持的算子不替换就跑不通没得商量复杂归一化子图LN、IN、GNTensorRT 大概率不自动融合一个个 launch kernel 性能差复杂的 attention 子图TensorRT 认不出来需要替换成它能优化的结构动态形状改静态TensorRT 对动态形状支持有限需要在 ONNX 层面简化修改节点属性/输入比如改 Reshape 的 shape只能在 ONNX 层面做怎么判断需不需要在 ONNX 层面做不要上来就手动融合所有东西按这个流程判断先直接用 TensorRT 跑看能不能跑通如果报错不支持某个算子→ 必须在 ONNX 层面处理如果能跑通 → 进入第二步看性能profiling 找瓶颈trtexec--onnxmodel.onnx--saveEnginemodel.engine--dumpProfile如果发现某一组小算子各自占一点时间加起来总耗时很大 → 说明 TensorRT 没有融合需要在 ONNX 层面手动融合如果常见的 ConvBN 已经被融合成一个节点了 → 不用管看 TensorRT 的优化日志加--verbose看融合前后的层数变化就能大概知道 TensorRT 做了多少融合经验法则常见模式ConvBN、ConvReLU→ 不用TensorRT 自动做复杂子图LayerNorm、attention、后处理→ 要做TensorRT 大概率不自动融合不支持的算子 → 必须做即使 TensorRT 会做ONNX 层面做也有价值。有些情况下即使 TensorRT 能自动融合在 ONNX 层面先做一遍也不是完全没意义价值说明跨平台兼容性融合后的 ONNX 模型在 onnxruntime、ncnn、OpenVINO 等其他后端上也能受益不只是 TensorRT模型可读性融合后的模型用 Netron 打开更简洁调试、排查问题更方便加快 TensorRT 构建速度TensorRT 做融合也需要时间模式匹配、图遍历ONNX 层面先做好可以缩短引擎构建时间避免版本差异不同版本的 TensorRT 融合能力不同在 ONNX 层面手动融合可以保证在所有版本上都有稳定的好性能减少中间张量的内存占用融合后中间张量不需要物化到内存所以即使是 ConvBN 这种 TensorRT 会自动融合的用 onnxsim 先跑一遍也没坏处——onnxsim 是自动的不需要你手动写代码成本很低。3.6 未知的、带参数的算子不好替换怎么办前面的案例都是已知的、无参数的或参数简单的算子替换。但实际工作中你可能会遇到未知的、带可训练参数的算子手动不好替换。这时候怎么办3.6.1 先分清是哪种情况情况例子是不是问题标准算子带参数Conv、Linear、BatchNorm❌ 不是问题ONNX 本来就支持参数自动导出为常量张量自定义算子带可训练参数自定义 attention 模块、自定义卷积变体、自定义激活函数⚠️ 是问题需要额外处理标准算子组合太碎LayerNorm 导出成 10 多个算子✅ 可以自动融合不是大问题你问的应该是第二类——自定义的、带可训练参数的算子。3.6.2 六种应对策略从优到劣策略一导出时就避免最优先成本最低最好的问题是不存在的问题。在模型设计或导出阶段就避免自定义算子用 ONNX 支持的标准算子重新实现自定义算子比如自定义的 attention可以用MatMul Softmax MatMul组合实现比如自定义的激活函数看看能不能用标准算子组合出来大部分自定义算子都能拆成标准算子的组合选高版本 opset 导出opset 越高支持的算子越多比如 LayerNorm 在 opset 17 才成为标准算子导出时尽量选最新的稳定 opset避免用 PyTorch 的自定义 autograd Functiontorch.autograd.Function导出时最容易出问题如果必须用确保写了 symbolic 方法3.4 的算子注册这是成本最低的方案能在导出阶段解决就不要拖到后面。策略二算子注册3.4 的方法最优雅如果 PyTorch 能识别这个算子比如 torchvision 里的、或者你自己写了 torch.ops 注册的只是 ONNX 导出器不认识那就用 3.4 的方法注册 symbolic 函数。带参数也没关系——参数weight、bias作为输入张量传过去导出后自动变成 ONNX 里的常量张量。这是最优雅的解决方案不需要手动改 ONNX导出时自动处理。策略三onnx_graphsurgeon 自动替换写模式匹配脚本如果已经导出成 ONNX 了而且子图模式是固定的可以写脚本自动识别和替换不用手动定位# 遍历所有节点找到符合某种模式的子图自动替换fornodeingraph.nodes:ifnode.opReduceMean:# 检查后面是不是跟着 Sub → Pow → ReduceMean → Add → Sqrt → Div...# 如果是就替换成单个 LayerNormalization...带参数也没关系——参数是常量张量替换时直接把旧的常量张量接过来就行。如果新算子需要不同格式的参数用gs.Constant创建新的常量张量把旧参数转换后放进去。本文的案例是手动定位的因为教学目的实际工作中可以写自动化的模式匹配脚本。策略四用 onnxsim 等工具自动简化很多常见的子图融合不需要你手动写脚本现成的工具自动就做了onnxsimONNX Simplifier自动消除冗余节点、常量折叠、算子融合ConvBN 融合、ConvRelu 融合、Identity 删除、Shape 相关节点简化……一行命令python -m onnxsim input.onnx output.onnxonnxruntime 的优化推理时自动做图优化TensorRT 的优化构建引擎时自动做算子融合和优化工作中拿到一个 ONNX 模型先跑一遍 onnxsim80% 的冗余和碎算子就自动处理了。剩下 onnxsim 处理不了的再手动用 gs 替换。策略五TensorRT Plugin最后手段最灵活如果以上方法都不行比如自定义算子太特殊无法拆成标准算子组合也无法等价替换就在TensorRT 里写自定义 Plugin把整个自定义算子包括参数、前向计算逻辑封装成一个 TensorRT PluginONNX 里只需要声明这个自定义算子节点用 gs 或者导出时注册TensorRT 推理时遇到这个节点就调用你写的 Plugin 来计算这是最灵活但也是最麻烦的方法需要写 CUDA kernel门槛高。一般到这一步说明这个算子确实很特殊。策略六重新训练万不得已成本最高如果自定义算子的结构和标准算子差异太大无法等价替换也无法写 Plugin或者写 Plugin 成本太高最后才考虑重新训练用标准算子重新实现一个功能类似的模块加载预训练权重然后微调fine-tune让新模块的输出尽量接近旧模块这是成本最高的方案需要数据、训练资源、时间而且可能有精度损失。一般不到万不得已不用。是算子设计不合理吗是 PyTorch 不合理吗都不是。这是整个行业的普遍问题叫研究-部署鸿沟research-deployment gapPyTorch 的设计目标是灵活的研究框架——支持任意自定义算子、任意计算图这是它的核心优势。如果限制只能用 ONNX 支持的算子那研究就没法做了新想法怎么验证ONNX 的设计目标是通用的中间表示——它不可能覆盖所有可能的自定义算子只能定义一套标准算子集覆盖 95% 的常见场景。部署框架TensorRT 等的设计目标是高性能推理——它们只优化标准算子自定义算子需要用户自己写 Plugin。这不是谁的错是三个不同目标的框架之间的天然矛盾。好的做法是研究阶段灵活设计随便用自定义算子快速验证想法部署阶段做算子工程化——把自定义算子改写成标准算子组合、或者写 Plugin、或者微调产品意识强的团队在模型设计阶段就考虑部署友好性尽量用标准算子避免太特殊的自定义结构总结本文围绕 ONNX 算子替换与子图融合的实践难点展开。首先介绍了定位子图的四种方法Netron 可视化最直观、gs 打印张量适合脚本探查、遍历节点可看清连接关系、按算子类型过滤适合大模型并给出先 Netron 看结构、再 gs 确认名字、最后写代码的标准流程。其次强调替换后必须做等价性验证分算子级、网络级、业务级三个层次并提供了可直接复用的验证代码模板同时区分了严格等价、近似等价与精度损失三种情况。最后深度解答了常见困惑ONNX 只是壳子、算子语义由官方文档定义、标准算子与自定义算子的区别、Identity 的工程用途、ONNX 与 TensorRT 分层融合的价值以及面对未知带参算子时的六种应对策略帮助读者系统掌握算子替换的完整方法论。本系列第一篇优雅地解决 onnx 算子不兼容 —— 算子注册本系列第二篇优雅地改造 onnx 计算图 —— 算子替换与子图融合上如果觉得有帮助欢迎点赞❤️、收藏❤️、关注❤️有问题可以在评论区交流。
