ONNX GraphSurgeon 节点删除实战重连图结构并用 cleanup 完成清理【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT导读在把 ONNX 模型导入 TensorRT 或做推理优化前模型中常常残留一些在推理期可有可无的节点例如多余的 Identity、Dropout 或自定义的占位节点。本教程以 TensorRT 仓库内 onnx-graphsurgeon 的06_removing_nodes示例为骨架完整讲解删除图中节点并重连剩余边的两步操作原理并结合源码剖析i()/o()便捷方法与cleanup()的清理机制。读完本文你将能够用几行代码安全地从任意 ONNX 图中摘除指定节点并导出结构整洁的新模型。一、示例要解决的问题本示例位于 tools/onnx-graphsurgeon/examples/06_removing_nodes共包含两个脚本与一份说明文档文件作用generate.py生成一个包含待删除节点的 ONNX 模型并保存为model.onnxremove.py从图中删除假节点重连剩余节点并保存为removed.onnxREADME.md说明节点删除的基本原理与运行步骤整个流程演示了 onnx-graphsurgeonONNX GraphSurgeon最核心的图编辑能力之一不是物理删除一个节点而是通过改写张量的生产者/消费者关系让数据流绕过该节点最后再借助cleanup()移除所有失去用途的孤儿节点与张量。二、先看 generate.py造出一个三节点小模型import onnx_graphsurgeon as gs import numpy as np import onnx # Inputs x gs.Variable(namex, dtypenp.float32, shape(1, 3, 224, 224)) # Intermediate tensors i0 gs.Variable(namei0) i1 gs.Variable(namei1) # Outputs y gs.Variable(namey, dtypenp.float32) nodes [ gs.Node(opIdentity, inputs[x], outputs[i0]), gs.Node(opFakeNodeToRemove, inputs[i0], outputs[i1]), gs.Node(opIdentity, inputs[i1], outputs[y]), ] graph gs.Graph(nodesnodes, inputs[x], outputs[y], ir_version10) model onnx.shape_inference.infer_shapes(gs.export_onnx(graph)) onnx.save(model, model.onnx)这段代码用 IR 层 API 从零构建了一张链式图x (1,3,224,224) | Identity (Node0) | i0 | FakeNodeToRemove (Node1) ← 目标把它删掉 | i1 | Identity (Node2) | y几个值得注意的细节张量分为两类带dtype与shape的输入x是Variable运行时才知道数值i0、i1只给了名字y只给了dtype这符合 tensor.py 中对Variable的定义——数值未知但可能携带数据类型与形状信息。gs.Graph(nodes..., inputs..., outputs..., ir_version10)直接指定图节点、图输入输出与 ONNX IR 版本导出时会写入模型头。导出前调用onnx.shape_inference.infer_shapes()做一次形状推断使保存的模型带有更完整的形状信息随后onnx.save落盘为model.onnx。运行方式python3 generate.py提示onnx-graphsurgeon 依赖onnx与numpy请先确认环境已安装。可通过python3 -m pip install onnx_graphsurgeon onnx numpy安装。三、节点删除的基本原理两步重连README 用一张链式图解释了删除的核心思想。假设初始结构为Node0 | Tensor0 | Node1 | Tensor1 | Node2 | Tensor2要把中间的Node1删除只需两步把Node0的输出从[Tensor0]改为[Tensor1]——让上游节点直接产出Node1的输出张量数据流就此绕过Node1把Node1的输出改为[]——摘掉Node1与Tensor1的产出关系使它从数据流中脱钩。修改后的图变为Node0 | Tensor1 Tensor0 | | Node2 Node1 | Tensor2可以看到此时Node1仍然残留在图的nodes列表中只是不再参与任何数据流。README 明确指出Note thatNode1remains in the graph. This is easily remedied by running thecleanup()function.Node1 仍留在图中运行cleanup()即可轻松解决。这个两步法之所以成立依赖 onnx-graphsurgeon IR 的双向同步机制Node.inputs/Node.outputs与Tensor.inputs/Tensor.outputs是SynchronizedList见 node.py 与 tensor.py对任意一侧的增删都会自动回写到所有关联的节点与张量。因此示例中改 Node0 的输出、清空 Node1 的输出两步就足以让整张图的关系保持一致无需手工同步维护另一侧的指针。四、再看 remove.py三步完成删除import onnx_graphsurgeon as gs import onnx graph gs.import_onnx(onnx.load(model.onnx)) fake_node [node for node in graph.nodes if node.op FakeNodeToRemove][0] # Get the input node of the fake node # Node provides i() and o() functions that can optionally be provided an index (default is 0) # These serve as convenience functions for the alternative, which would be to fetch the input/output # tensor first, then fetch the input/output node of the tensor. # For example, node.i() is equivalent to node.inputs[0].inputs[0] inp_node fake_node.i() # Reconnect the input node to the output tensors of the fake node, so that the first identity # node in the example graph now skips over the fake node. inp_node.outputs fake_node.outputs fake_node.outputs.clear() # Remove the fake node from the graph completely graph.cleanup() model onnx.shape_inference.infer_shapes(gs.export_onnx(graph)) onnx.save(model, removed.onnx)逐行拆解导入图gs.import_onnx(onnx.load(model.onnx))把 ONNX 模型读入 onnx-graphsurgeon IR。导入接口定义在 base_importer.py高层入口是gs.import_onnx。定位目标节点用列表推导式按op名筛选出FakeNodeToRemove。在真实场景中通常会按node.op、node.name或输入输出张量名组合定位待删节点。取上游节点fake_node.i()返回Node1的输入生产者。在 node.py 中i()的定义是self.inputs[tensor_idx].inputs[producer_idx]默认等价于node.inputs[0].inputs[0]对应地o()定义是self.outputs[tensor_idx].outputs[consumer_idx]用于取下游消费者。参数含义为i(tensor_idx, producer_idx)tensor_idx指定取本节点的第几个输入张量producer_idx指定该张量的第几个生产者张量一般只有一个生产者o(consumer_idx, tensor_idx)consumer_idx指定取本节点第几个输出张量的消费者tensor_idx指定输出张量索引。两步重连对应第三章的核心操作inp_node.outputs fake_node.outputs把上游 Identity 的输出直接接到Node1的输出张量i1上fake_node.outputs.clear()清空Node1的输出令其从数据流中脱钩。彻底清理graph.cleanup()移除不再服务于任何图输出的节点与张量详见下一节。导出落盘与生成流程对称先infer_shapes再做形状推断最后onnx.save保存为removed.onnx。运行方式python3 remove.py五、cleanup() 源码级解析它到底清理了什么cleanup()定义在 graph.py完整签名为def cleanup( self, remove_unused_node_outputsFalse, recurse_subgraphsTrue, remove_unused_graph_inputsFalse, recurse_functionsTrue, ):其核心判定标准是一个节点或张量只要不贡献给任何图输出就被视为未使用。实现上从图输出出发做反向遍历_get_used_node_ids内部维护used_tensors队列沿used_tensor.inputs逐级回溯生产者凡是命中的节点才被保留未命中的节点会被清空inputs/outputs并从graph.nodes移除。几个参数的作用参数默认值说明remove_unused_node_outputsFalse是否一并移除节点上悬空的输出张量没有消费者的张量永远不会删除图输出张量也不会删除名字为空的可选占位输出recurse_subgraphsTrue是否递归清理子图如If、Loop控制流节点内嵌的图remove_unused_graph_inputsFalse是否移除未使用的图输入recurse_functionsTrue是否同时清理图的 Local Functions本地函数文档还特别强调了两点边界cleanup()永远不会修改图输出张量同时它也会剔除图输入张量的生产者以及图输出张量的、不在图中的消费者保证图边界干净。回到本示例两步重连后Node1不再产生任何被消费的数据cleanup()遍历时它不在used_node_ids中于是连同它产出的孤立张量i0以及已无消费者的i1旧引用一起被剔除最终图中只剩Identity - Identity两条边即上面第二张图所示的结构。六、真实场景什么时候需要摘节点虽然示例用的是虚构的FakeNodeToRemove但同样的两步重连模式可以直接套用到日常 ONNX 模型优化中典型场景包括删除恒等/冗余节点训练模型中的Identity、Dropout推理时无实际作用等删除自定义占位节点某些工具导出的模型中带Constant占位、调试用节点推理期应当摘除配合子图替换/常量折叠使用例如先graph.fold_constants()把常量子图折叠成Constant张量再配合cleanup()清掉折叠后残留的空转节点fold_constants的注释明确建议折叠后调用cleanup()移除悬挂节点删除不满足 TensorRT 算子约束的节点为后续导入 TensorRT 前精简图结构。定位待删节点时除了按op名称也可以组合node.name、node.inputs/node.outputs的张量名进行过滤删除后建议继续调用graph.toposort()对节点重新拓扑排序确保导出顺序正确。七、验证结果与延伸阅读运行完两个脚本后工作目录会出现两个文件model.onnx包含FakeNodeToRemove的原始模型removed.onnx删除并清理后的模型。可以用 Netron 打开两者对比本示例 README 中的配图即 Netron 可视化结果也可以再用gs.import_onnx(onnx.load(removed.onnx))检查graph.nodes中是否只剩下两个Identity。如果想继续深入 onnx-graphsurgeon 的图编辑能力推荐按顺序阅读同目录下的其他示例04_modifying_a_model修改图中节点的通用手法是删除节点之外最常用的编辑操作05_folding_constants常量折叠常与cleanup()搭配使用08_replacing_a_subgraph用自定义子图替换图中一段结构是比单节点删除更进阶的重连操作。底层 IR 细节可查看 node.pyi()/o()与同步列表、graph.pycleanup()/toposort()/tensors()以及 README.mdIR 总体设计说明。【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考