onnx-graphsurgeon 的 Graph 类完全指南掌握 TensorRT 生态中 ONNX 计算图的建模、遍历与优化【免费下载链接】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/TensorRTonnx-graphsurgeon 是 NVIDIA TensorRT 开源仓库中随附的 ONNX 模型创建与修改库而Graph类是其中间表示IR的核心容器承载着 ONNX 计算图的全部节点、张量、函数与元信息。本文以 docs/ir/graph.rst 文档所指向的onnx_graphsurgeon.Graph类为主线结合 graph.py 的完整源码实现系统讲解 Graph 的构造、注册机制、节点增删、拓扑排序、常量折叠与子图处理等核心能力帮助你掌握在 TensorRT 导入 ONNX 模型前对计算图进行精细化改造的完整方法论。Graph 在 onnx-graphsurgeon IR 中的定位onnx-graphsurgeon 由三大部分组成Importers导入器将 ONNX 模型导入 IR、IR中间表示修改与新建计算图的地方、Exporters导出器将 IR 写回 ONNX。其中 IR 由三个互相引用的组件构成Tensortensor.py张量分为值已知的Constant与运行期才确定值的Variable两个子类Nodenode.py节点表示一个算子操作消费零个或多个 Tensor产出零个或多个 TensorGraphgraph.py计算图持有节点列表nodes、图输入inputs、图输出outputs是用户操作的主入口。Graph类的类注释只有一句话——“Represents a graph containing nodes and tensors”但它的完整 API 覆盖面非常广。日常最典型的用法是先通过导入器把 ONNX 模型变成 Graphimport onnx import onnx_graphsurgeon as gs graph gs.import_onnx(onnx.load(model.onnx)) # ... 对 graph 进行各种修改 ... onnx.save(gs.export_onnx(graph), model_modified.onnx)从源码结构看导入与导出分别由onnx_graphsurgeon/importers/与onnx_graphsurgeon/exporters/两个目录中的实现完成Graph处于两者的正中间是连接“ONNX 模型文件”与“可编辑的 Python 对象模型”的桥梁。构造一个 Graph构造参数与默认值Graph.__init__graph.py支持通过显式传入节点与张量列表来从零构建计算图也可由导入器自动填充。各参数及默认值如下参数类型默认值说明nodesSequence[Node][]图中的节点列表inputsSequence[Tensor][]图输入张量列表outputsSequence[Tensor][]图输出张量列表namestronnx_graphsurgeon_graph图名称doc_stringstr图文档字符串opsetint11Graph.DEFAULT_OPSET导出该图时使用的 ONNX opsetimport_domains序列None导入的自定义 domain 列表producer_namestr生成模型所用工具的名称producer_versionstr生成工具的版本functionsSequence[Function][]图中的本地函数ONNX Function列表ir_versionintNone导出时使用的 ONNX IR 版本几个值得注意的实现细节inputs、outputs在赋值时会被强制转换为普通list见__setattr__避免它们被包装成SynchronizedList_functions属性加了下划线前缀注释明确说明在 ONNX 中同一份函数列表会在模型内所有 Graph 与 Function 之间共享用户应通过graph.functions属性访问构造时还会调用_merge_subgraph_functions()把子图中携带的函数列表合并进父图保证整个模型共享同一份函数列表。import onnx_graphsurgeon as gs from onnx_graphsurgeon.ir.tensor import Variable x Variable(namex, dtypefloat32, shape[1, 3, 224, 224]) y gs.Variable(namey, dtypefloat32) node gs.Node(opRelu, namerelu_0, inputs[x], outputs[y]) graph gs.Graph(nodes[node], inputs[x], outputs[y], namemy_graph, opset13)用register扩展 Graph 方法面向 opset 的自定义算子注册Graph.registergraph.py是一个静态装饰器允许把自定义函数注册为Graph的成员方法之后即可像普通方法一样调用。核心机制是不传opsets时函数注册到GLOBAL_FUNC_MAP对所有 opset 生效传入 opset 列表时函数按 opset 分别注册到OPSET_FUNC_MAP同名函数可为不同 opset 各注册一份同一 opset 下重复注册会覆盖旧实现若注册名与 Graph 现有属性/方法重名会输出警告“该函数永远不会被调用”。gs.Graph.register() def add(self, a, b): return self.layer(opAdd, inputs[a, b], outputs[add_out_gs]) graph.add(a, b)__getattr__graph.py中的查找优先级为当前 opset 特定注册 全局注册 本地 Function若同名存在多个候选会打印一次警告并选择第一个候选。若目标函数未在当前 opset 注册会给出错误日志提示该函数实际注册在哪些 opset 上。用layer()创建节点灵活多样的输入输出类型Graph.layergraph.py是“创建节点并加入图”的统一入口它创建节点、加入self.nodes并返回该节点的输出张量列表。其核心设计在于inputs/outputs列表可混合以下四种类型由process_io统一归一化传入元素类型处理方式Tensor原样使用需保证张量名全局唯一str以该字符串为前缀自动生成唯一命名的Variablenumpy.ndarray以onnx_graphsurgeon_constant为前缀生成唯一命名的Constantlist/tuple/numbers.Number生成 1D 常量张量前缀为onnx_graphsurgeon_lst_constant含浮点数时 dtype 为np.float32否则为np.int64名称唯一性由_generate_name保证在既有名称集合中循环追加索引直到不冲突为止。节点自身名称默认以onnx_graphsurgeon_node前缀自动生成。未识别的元素类型会触发 critical 错误并提示“是否忘记用*解包列表”。# 传入字符串自动建 Variable、传入 numpy 数组自动建 Constant outs graph.layer( opConv, inputs[input, np.random.rand(64, 3, 3, 3).astype(np.float32)], outputs[conv_out], attrs{strides: [1, 1], pads: [1, 1, 1, 1]}, )图内张量与节点导航tensors()、subgraphs()、node_ids()tensors()构建张量名到张量的映射Graph.tensorsgraph.py遍历图内所有节点的输入输出以及图 I/O返回一个OrderedDict[str, Tensor]保证张量顺序与图中节点顺序一致——因此若图已拓扑排序该张量映射也是有序的。空张量is_empty()即名字为用于表示省略的可选输入输出会被忽略。若图中存在同名但不同对象的张量默认打印警告check_duplicatesTrue时则直接抛出OnnxGraphSurgeonException。subgraphs()遍历控制流子图Graph.subgraphsgraph.py以生成器方式迭代图中包含的所有子图。子图通常出现在 ONNX 控制流节点如If、Loop的属性中——这正是 node.py 中Node.subgraphs()的逻辑遍历节点属性凡是属性值为Graph实例的都作为子图产出recursiveTrue时还会继续深入子图节点。Graph.subgraphs(recursiveTrue)则递归搜索所有嵌套层级。node_ids()为节点分配临时唯一整数 IDGraph.node_idsgraph.py返回一个上下文管理器NodeIDAdder在进入时按节点在列表中的索引为每个节点打上id属性退出时删除。这一机制让“列表索引”与“节点对象”解耦——即使同一对象在列表中出现多次也能被区分。内部算法如cleanup、toposort大量依赖它同时建议所有内部代码通过_get_node_id获取节点 ID。with graph.node_ids(): assert graph.nodes[0].id ! graph.nodes[1].idcleanup()剔除无用节点与张量Graph.cleanupgraph.py是图优化中最常用的“瘦身”操作任何不贡献于图输出的节点与张量都被视为无用并删除。算法从图输出出发沿张量-节点反向遍历_get_used_node_ids收集所有被使用的节点集合注意此过程只考虑本图的局部张量来自外层图的“外来张量”_foreign_tensors不会被误删。关键参数与语义参数默认值说明remove_unused_node_outputsFalse是否删除节点的无用输出张量悬挂张量非空、无消费者且非图输出空张量省略的可选输出永远不会被删除recurse_subgraphsTrue是否递归清理子图remove_unused_graph_inputsFalse是否删除未被使用的图输入recurse_functionsTrue是否同时清理本地函数的节点实现上还有两个细节图输入的产生者节点会被先行清除以保证used_node_ids计算正确函数与子图的清理各自只执行一次避免无限递归。注意cleanup 永远不会修改图输出张量。graph.cleanup(remove_unused_node_outputsTrue, remove_unused_graph_inputsTrue)toposort()原地拓扑排序Graph.toposortgraph.py将图原地按拓扑序排列保证每个节点的输入都先于该节点被计算。这是调用fold_constants等依赖有序遍历功能的前置条件。mode参数控制排序范围full默认同时排序节点列表与函数列表nodes仅排序节点functions仅排序函数。排序算法为每个节点计算“层级”hierarchy level节点的层级 其最高层输入节点的层级 1输入节点层级为 0。计算时会考虑节点输入张量的本地产生者、子图使用的外来张量以及函数间的调用关系。实现中还使用visited集合检测环——若发现环会以 critical 日志提示“图中是否存在重名张量”。对菱形依赖图代码注释特别指出“先查已缓存层级再查 visited”以避免误报假环。graph.toposort()fold_constants()常量折叠的完整流水线Graph.fold_constantsgraph.py在图中原地折叠常量把“只依赖常量的子表达式”预先计算成常量。源码注释明确了两条前置约束调用前图必须已拓扑排序图必须可导出为 ONNX 且可被 ONNX-Runtime 求值且环境中必须安装 onnxruntime。折叠后产生的悬挂节点不会自动删除需要再调用cleanup()。核心参数一览参数默认值说明fold_shapesTrue是否折叠Shape节点要求图中已推断形状且仅折叠静态形状recurse_subgraphsTrue是否递归折叠子图中的常量partitioningNone图分区策略None不分区一次推理失败则整体不折叠、basic分区某分区失败不影响其他分区、recursive递归再分区error_okTrue是否抑制推理错误为False时错误会被重新抛出flatten_subgraphsTrue条件为常量的If节点可被扁平化融入父图size_thresholdNone折叠的最大字节阈值超过该值的张量不折叠None表示不限制。对Tile/Expand等可能产生超大常量的算子尤为有用should_exclude_node恒返回False回调函数返回True的节点不参与折叠其消费者也不会被折叠recurse_functionsTrue是否折叠本地函数中的常量五趟流水线源码注释详细描述了fold_constants内部执行的 5 个 passPass 1——常量节点降级把Constant算子节点降级为Constant张量支持value、value_float、value_floats、value_int、value_ints等属性Pass 2——形状张量 Cast 消除针对Cast(int→float) → 中间算子与 float 常量运算 → Cast(→int)这种对 TensorRT 形状张量不友好的模式做消除。候选中间算子包括Add/Sub/Mul/Div/Max/Min/Equal/Greater/Less/Concat并要求输入 Cast 均来自同一源类型、输出 Cast 目标类型一致且等于源类型Pass 3——寻找常量张量的所有后代在图的克隆副本上沿拓扑序扩展“仅依赖常量”的可折叠节点集合QuantizeLinear、DequantizeLinear、DynamicQuantizeLinear被排除在可折叠范围之外NO_FOLD_OPSPass 4——Shape 折叠依次执行fold_shape_gather、fold_shape_slice、fold_shape三个子例程分别把Shape→Gather、Shape→Slice、Shape折叠为静态形状常量。仅当输入形状静态已知且各索引/切片参数为常量时才可折叠axes必须为 0形状张量仅支持第 0 维切片。源码注释特别提示了相位顺序问题三个子例程的执行顺序会影响最终折叠量Pass 5——ONNX-Runtime 求值将可折叠张量导出为 ONNX 子图用onnxruntime.InferenceSessionCPUExecutionProvider求值再把结果替换回原图。启用partitioning时按输出节点逐个分区推理失败则按recursive逐节点剥离重试。大张量保护折叠完成后若产生超过 1 MiB1 20字节的常量会以LogMode.ONCE模式输出一次警告提示模型可能因折叠大张量而膨胀建议设置size_threshold显式设置了阈值时超限张量会被跳过并输出 debug 日志。此外条件为常量张量的If节点会被扁平化then_branch/else_branch子图的局部张量会追加_subg_{index}_{name}后缀避免与父图张量冲突子图输出重绑定到If节点的输出然后子图节点内联进父图。使用示例graph.fold_constants() # 默认行为折叠形状 递归子图 graph.cleanup() # 移除折叠后留下的悬挂节点 # 带分区与大小阈值避免大张量撑爆模型 graph.fold_constants( partitioningrecursive, size_threshold1024 * 1024, should_exclude_nodelambda node: node.op CustomOp, )张量的就地类型转换to_constant/to_variable尽管这两个方法定义在Tensor基类上却是图改造中高频使用的工具。tensor.py 中Tensor.to_constant(values, data_locationNone, export_dtypeNone)通过直接修改self.__class__把张量原地变为Constant所有消费者/产生者都会看到更新to_variable(dtype, shape)则反向把常量变成Variable。在fold_constants的实现中折叠结果正是通过tensor.to_constant(values)写回原图并清除输入来实现的。copy()复制图而不深拷贝权重Graph.copygraph.py复制整个图所有节点与张量都会被拷贝但不会深拷贝权重或属性唯一的例外是Graph类型的属性会调用其自身的copy方法。tensor_map参数用于提供外层图的张量映射——当单独克隆一个子图时需要通过tensors()拿到全部所需张量包括外层图产生的合并优先级为“外层图已拷贝版本 本地张量拷贝”即本地产生的张量拥有最高优先级。fold_constants内部正是先用self.copy()得到克隆图用于 ONNX-Runtime 求值再回写原图。图的字符串表示与相等性__str__graph.py输出形如Graph {name} (Opset {opset})的可读描述包含 Local Functions、Inputs、Nodes、Outputs 四段信息__eq__graph.py判定两个图相等要求opset、import_domains、ir_version、函数集合按unique_id比较以及 inputs/outputs/nodes 序列均相等。组合实战一次完整的图改造流程结合上述 API一个典型的“导入-修改-导出”工作流如下先toposort()保证有序再fold_constants()cleanup()完成常量折叠与瘦身最后导出回 ONNX 供 TensorRT 使用import onnx import onnx_graphsurgeon as gs graph gs.import_onnx(onnx.load(model.onnx)) # 1. 拓扑排序fold_constants 的前置条件 graph.toposort() # 2. 常量折叠 清理悬挂节点 graph.fold_constants(partitioningbasic) graph.cleanup(remove_unused_node_outputsTrue) # 3. 删除指定算子以 Dropout 为例 graph.nodes [n for n in graph.nodes if n.op ! Dropout] # 4. 替换图输出 graph.outputs [graph.tensors()[new_output]] # 5. 导出并保存 onnx.save(gs.export_onnx(graph), model_optimized.onnx)进一步阅读IR 文档目录Graph、Node、Tensor 三部分的 Sphinx 文档索引Graph 类完整源码本文所有 API 与内部算法的第一手依据Node 类源码 与 Tensor 类源码IR 的另外两个组成组件项目 README包含安装方式python3 -m pip install onnx_graphsurgeon、导入器/IR/导出器三大组件总览examples 示例目录若干常见用法的可直接运行示例。在 TensorRT 的模型转换流程中Graph类是与 ONNX 模型交互的核心入口——无论是裁剪无关算子、折叠常量、修正形状张量还是为不支持的算子改写子图几乎所有图级操作都以它为操作对象展开理解其 API 与内部算法是高效使用 onnx-graphsurgeon 的前提。【免费下载链接】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），仅供参考 onnx-graphsurgeon 的 Graph 类完全指南：掌握 TensorRT 生态中 ONNX 计算图的建模、遍历与优化
onnx-graphsurgeon 的 Graph 类完全指南掌握 TensorRT 生态中 ONNX 计算图的建模、遍历与优化【免费下载链接】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/TensorRTonnx-graphsurgeon 是 NVIDIA TensorRT 开源仓库中随附的 ONNX 模型创建与修改库而Graph类是其中间表示IR的核心容器承载着 ONNX 计算图的全部节点、张量、函数与元信息。本文以 docs/ir/graph.rst 文档所指向的onnx_graphsurgeon.Graph类为主线结合 graph.py 的完整源码实现系统讲解 Graph 的构造、注册机制、节点增删、拓扑排序、常量折叠与子图处理等核心能力帮助你掌握在 TensorRT 导入 ONNX 模型前对计算图进行精细化改造的完整方法论。Graph 在 onnx-graphsurgeon IR 中的定位onnx-graphsurgeon 由三大部分组成Importers导入器将 ONNX 模型导入 IR、IR中间表示修改与新建计算图的地方、Exporters导出器将 IR 写回 ONNX。其中 IR 由三个互相引用的组件构成Tensortensor.py张量分为值已知的Constant与运行期才确定值的Variable两个子类Nodenode.py节点表示一个算子操作消费零个或多个 Tensor产出零个或多个 TensorGraphgraph.py计算图持有节点列表nodes、图输入inputs、图输出outputs是用户操作的主入口。Graph类的类注释只有一句话——“Represents a graph containing nodes and tensors”但它的完整 API 覆盖面非常广。日常最典型的用法是先通过导入器把 ONNX 模型变成 Graphimport onnx import onnx_graphsurgeon as gs graph gs.import_onnx(onnx.load(model.onnx)) # ... 对 graph 进行各种修改 ... onnx.save(gs.export_onnx(graph), model_modified.onnx)从源码结构看导入与导出分别由onnx_graphsurgeon/importers/与onnx_graphsurgeon/exporters/两个目录中的实现完成Graph处于两者的正中间是连接“ONNX 模型文件”与“可编辑的 Python 对象模型”的桥梁。构造一个 Graph构造参数与默认值Graph.__init__graph.py支持通过显式传入节点与张量列表来从零构建计算图也可由导入器自动填充。各参数及默认值如下参数类型默认值说明nodesSequence[Node][]图中的节点列表inputsSequence[Tensor][]图输入张量列表outputsSequence[Tensor][]图输出张量列表namestronnx_graphsurgeon_graph图名称doc_stringstr图文档字符串opsetint11Graph.DEFAULT_OPSET导出该图时使用的 ONNX opsetimport_domains序列None导入的自定义 domain 列表producer_namestr生成模型所用工具的名称producer_versionstr生成工具的版本functionsSequence[Function][]图中的本地函数ONNX Function列表ir_versionintNone导出时使用的 ONNX IR 版本几个值得注意的实现细节inputs、outputs在赋值时会被强制转换为普通list见__setattr__避免它们被包装成SynchronizedList_functions属性加了下划线前缀注释明确说明在 ONNX 中同一份函数列表会在模型内所有 Graph 与 Function 之间共享用户应通过graph.functions属性访问构造时还会调用_merge_subgraph_functions()把子图中携带的函数列表合并进父图保证整个模型共享同一份函数列表。import onnx_graphsurgeon as gs from onnx_graphsurgeon.ir.tensor import Variable x Variable(namex, dtypefloat32, shape[1, 3, 224, 224]) y gs.Variable(namey, dtypefloat32) node gs.Node(opRelu, namerelu_0, inputs[x], outputs[y]) graph gs.Graph(nodes[node], inputs[x], outputs[y], namemy_graph, opset13)用register扩展 Graph 方法面向 opset 的自定义算子注册Graph.registergraph.py是一个静态装饰器允许把自定义函数注册为Graph的成员方法之后即可像普通方法一样调用。核心机制是不传opsets时函数注册到GLOBAL_FUNC_MAP对所有 opset 生效传入 opset 列表时函数按 opset 分别注册到OPSET_FUNC_MAP同名函数可为不同 opset 各注册一份同一 opset 下重复注册会覆盖旧实现若注册名与 Graph 现有属性/方法重名会输出警告“该函数永远不会被调用”。gs.Graph.register() def add(self, a, b): return self.layer(opAdd, inputs[a, b], outputs[add_out_gs]) graph.add(a, b)__getattr__graph.py中的查找优先级为当前 opset 特定注册 全局注册 本地 Function若同名存在多个候选会打印一次警告并选择第一个候选。若目标函数未在当前 opset 注册会给出错误日志提示该函数实际注册在哪些 opset 上。用layer()创建节点灵活多样的输入输出类型Graph.layergraph.py是“创建节点并加入图”的统一入口它创建节点、加入self.nodes并返回该节点的输出张量列表。其核心设计在于inputs/outputs列表可混合以下四种类型由process_io统一归一化传入元素类型处理方式Tensor原样使用需保证张量名全局唯一str以该字符串为前缀自动生成唯一命名的Variablenumpy.ndarray以onnx_graphsurgeon_constant为前缀生成唯一命名的Constantlist/tuple/numbers.Number生成 1D 常量张量前缀为onnx_graphsurgeon_lst_constant含浮点数时 dtype 为np.float32否则为np.int64名称唯一性由_generate_name保证在既有名称集合中循环追加索引直到不冲突为止。节点自身名称默认以onnx_graphsurgeon_node前缀自动生成。未识别的元素类型会触发 critical 错误并提示“是否忘记用*解包列表”。# 传入字符串自动建 Variable、传入 numpy 数组自动建 Constant outs graph.layer( opConv, inputs[input, np.random.rand(64, 3, 3, 3).astype(np.float32)], outputs[conv_out], attrs{strides: [1, 1], pads: [1, 1, 1, 1]}, )图内张量与节点导航tensors()、subgraphs()、node_ids()tensors()构建张量名到张量的映射Graph.tensorsgraph.py遍历图内所有节点的输入输出以及图 I/O返回一个OrderedDict[str, Tensor]保证张量顺序与图中节点顺序一致——因此若图已拓扑排序该张量映射也是有序的。空张量is_empty()即名字为用于表示省略的可选输入输出会被忽略。若图中存在同名但不同对象的张量默认打印警告check_duplicatesTrue时则直接抛出OnnxGraphSurgeonException。subgraphs()遍历控制流子图Graph.subgraphsgraph.py以生成器方式迭代图中包含的所有子图。子图通常出现在 ONNX 控制流节点如If、Loop的属性中——这正是 node.py 中Node.subgraphs()的逻辑遍历节点属性凡是属性值为Graph实例的都作为子图产出recursiveTrue时还会继续深入子图节点。Graph.subgraphs(recursiveTrue)则递归搜索所有嵌套层级。node_ids()为节点分配临时唯一整数 IDGraph.node_idsgraph.py返回一个上下文管理器NodeIDAdder在进入时按节点在列表中的索引为每个节点打上id属性退出时删除。这一机制让“列表索引”与“节点对象”解耦——即使同一对象在列表中出现多次也能被区分。内部算法如cleanup、toposort大量依赖它同时建议所有内部代码通过_get_node_id获取节点 ID。with graph.node_ids(): assert graph.nodes[0].id ! graph.nodes[1].idcleanup()剔除无用节点与张量Graph.cleanupgraph.py是图优化中最常用的“瘦身”操作任何不贡献于图输出的节点与张量都被视为无用并删除。算法从图输出出发沿张量-节点反向遍历_get_used_node_ids收集所有被使用的节点集合注意此过程只考虑本图的局部张量来自外层图的“外来张量”_foreign_tensors不会被误删。关键参数与语义参数默认值说明remove_unused_node_outputsFalse是否删除节点的无用输出张量悬挂张量非空、无消费者且非图输出空张量省略的可选输出永远不会被删除recurse_subgraphsTrue是否递归清理子图remove_unused_graph_inputsFalse是否删除未被使用的图输入recurse_functionsTrue是否同时清理本地函数的节点实现上还有两个细节图输入的产生者节点会被先行清除以保证used_node_ids计算正确函数与子图的清理各自只执行一次避免无限递归。注意cleanup 永远不会修改图输出张量。graph.cleanup(remove_unused_node_outputsTrue, remove_unused_graph_inputsTrue)toposort()原地拓扑排序Graph.toposortgraph.py将图原地按拓扑序排列保证每个节点的输入都先于该节点被计算。这是调用fold_constants等依赖有序遍历功能的前置条件。mode参数控制排序范围full默认同时排序节点列表与函数列表nodes仅排序节点functions仅排序函数。排序算法为每个节点计算“层级”hierarchy level节点的层级 其最高层输入节点的层级 1输入节点层级为 0。计算时会考虑节点输入张量的本地产生者、子图使用的外来张量以及函数间的调用关系。实现中还使用visited集合检测环——若发现环会以 critical 日志提示“图中是否存在重名张量”。对菱形依赖图代码注释特别指出“先查已缓存层级再查 visited”以避免误报假环。graph.toposort()fold_constants()常量折叠的完整流水线Graph.fold_constantsgraph.py在图中原地折叠常量把“只依赖常量的子表达式”预先计算成常量。源码注释明确了两条前置约束调用前图必须已拓扑排序图必须可导出为 ONNX 且可被 ONNX-Runtime 求值且环境中必须安装 onnxruntime。折叠后产生的悬挂节点不会自动删除需要再调用cleanup()。核心参数一览参数默认值说明fold_shapesTrue是否折叠Shape节点要求图中已推断形状且仅折叠静态形状recurse_subgraphsTrue是否递归折叠子图中的常量partitioningNone图分区策略None不分区一次推理失败则整体不折叠、basic分区某分区失败不影响其他分区、recursive递归再分区error_okTrue是否抑制推理错误为False时错误会被重新抛出flatten_subgraphsTrue条件为常量的If节点可被扁平化融入父图size_thresholdNone折叠的最大字节阈值超过该值的张量不折叠None表示不限制。对Tile/Expand等可能产生超大常量的算子尤为有用should_exclude_node恒返回False回调函数返回True的节点不参与折叠其消费者也不会被折叠recurse_functionsTrue是否折叠本地函数中的常量五趟流水线源码注释详细描述了fold_constants内部执行的 5 个 passPass 1——常量节点降级把Constant算子节点降级为Constant张量支持value、value_float、value_floats、value_int、value_ints等属性Pass 2——形状张量 Cast 消除针对Cast(int→float) → 中间算子与 float 常量运算 → Cast(→int)这种对 TensorRT 形状张量不友好的模式做消除。候选中间算子包括Add/Sub/Mul/Div/Max/Min/Equal/Greater/Less/Concat并要求输入 Cast 均来自同一源类型、输出 Cast 目标类型一致且等于源类型Pass 3——寻找常量张量的所有后代在图的克隆副本上沿拓扑序扩展“仅依赖常量”的可折叠节点集合QuantizeLinear、DequantizeLinear、DynamicQuantizeLinear被排除在可折叠范围之外NO_FOLD_OPSPass 4——Shape 折叠依次执行fold_shape_gather、fold_shape_slice、fold_shape三个子例程分别把Shape→Gather、Shape→Slice、Shape折叠为静态形状常量。仅当输入形状静态已知且各索引/切片参数为常量时才可折叠axes必须为 0形状张量仅支持第 0 维切片。源码注释特别提示了相位顺序问题三个子例程的执行顺序会影响最终折叠量Pass 5——ONNX-Runtime 求值将可折叠张量导出为 ONNX 子图用onnxruntime.InferenceSessionCPUExecutionProvider求值再把结果替换回原图。启用partitioning时按输出节点逐个分区推理失败则按recursive逐节点剥离重试。大张量保护折叠完成后若产生超过 1 MiB1 20字节的常量会以LogMode.ONCE模式输出一次警告提示模型可能因折叠大张量而膨胀建议设置size_threshold显式设置了阈值时超限张量会被跳过并输出 debug 日志。此外条件为常量张量的If节点会被扁平化then_branch/else_branch子图的局部张量会追加_subg_{index}_{name}后缀避免与父图张量冲突子图输出重绑定到If节点的输出然后子图节点内联进父图。使用示例graph.fold_constants() # 默认行为折叠形状 递归子图 graph.cleanup() # 移除折叠后留下的悬挂节点 # 带分区与大小阈值避免大张量撑爆模型 graph.fold_constants( partitioningrecursive, size_threshold1024 * 1024, should_exclude_nodelambda node: node.op CustomOp, )张量的就地类型转换to_constant/to_variable尽管这两个方法定义在Tensor基类上却是图改造中高频使用的工具。tensor.py 中Tensor.to_constant(values, data_locationNone, export_dtypeNone)通过直接修改self.__class__把张量原地变为Constant所有消费者/产生者都会看到更新to_variable(dtype, shape)则反向把常量变成Variable。在fold_constants的实现中折叠结果正是通过tensor.to_constant(values)写回原图并清除输入来实现的。copy()复制图而不深拷贝权重Graph.copygraph.py复制整个图所有节点与张量都会被拷贝但不会深拷贝权重或属性唯一的例外是Graph类型的属性会调用其自身的copy方法。tensor_map参数用于提供外层图的张量映射——当单独克隆一个子图时需要通过tensors()拿到全部所需张量包括外层图产生的合并优先级为“外层图已拷贝版本 本地张量拷贝”即本地产生的张量拥有最高优先级。fold_constants内部正是先用self.copy()得到克隆图用于 ONNX-Runtime 求值再回写原图。图的字符串表示与相等性__str__graph.py输出形如Graph {name} (Opset {opset})的可读描述包含 Local Functions、Inputs、Nodes、Outputs 四段信息__eq__graph.py判定两个图相等要求opset、import_domains、ir_version、函数集合按unique_id比较以及 inputs/outputs/nodes 序列均相等。组合实战一次完整的图改造流程结合上述 API一个典型的“导入-修改-导出”工作流如下先toposort()保证有序再fold_constants()cleanup()完成常量折叠与瘦身最后导出回 ONNX 供 TensorRT 使用import onnx import onnx_graphsurgeon as gs graph gs.import_onnx(onnx.load(model.onnx)) # 1. 拓扑排序fold_constants 的前置条件 graph.toposort() # 2. 常量折叠 清理悬挂节点 graph.fold_constants(partitioningbasic) graph.cleanup(remove_unused_node_outputsTrue) # 3. 删除指定算子以 Dropout 为例 graph.nodes [n for n in graph.nodes if n.op ! Dropout] # 4. 替换图输出 graph.outputs [graph.tensors()[new_output]] # 5. 导出并保存 onnx.save(gs.export_onnx(graph), model_optimized.onnx)进一步阅读IR 文档目录Graph、Node、Tensor 三部分的 Sphinx 文档索引Graph 类完整源码本文所有 API 与内部算法的第一手依据Node 类源码 与 Tensor 类源码IR 的另外两个组成组件项目 README包含安装方式python3 -m pip install onnx_graphsurgeon、导入器/IR/导出器三大组件总览examples 示例目录若干常见用法的可直接运行示例。在 TensorRT 的模型转换流程中Graph类是与 ONNX 模型交互的核心入口——无论是裁剪无关算子、折叠常量、修正形状张量还是为不支持的算子改写子图几乎所有图级操作都以它为操作对象展开理解其 API 与内部算法是高效使用 onnx-graphsurgeon 的前提。【免费下载链接】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），仅供参考