Binary search to locate the faulty subgraph

Author: TelGome

graph_net/torch/decompose_util.py

@@ -215,13 +215,22 @@ def _get_submodule_inputs_and_outputs(
     )
     node_list = list(gm.graph.nodes)
 
+    def _hashable(obj):
+        if isinstance(obj, slice):
+            return ("__slice__", obj.start, obj.stop, obj.step)
+        elif isinstance(obj, (list, tuple)):
+            return tuple(_hashable(x) for x in obj)
+        else:
+            return obj
+
     def get_related_node(node):
         for arg in node.args:
             if isinstance(arg, tuple):
-                yield from arg
+                for x in arg:
+                    yield _hashable(x)
             else:
-                yield arg
-        yield node
+                yield _hashable(arg)
+        yield _hashable(node)
 
     for node in node_list[0:start_node_idx]:
         for related_node in get_related_node(node):

graph_net/torch/naive_graph_decomposer.py

@@ -7,6 +7,23 @@
 from graph_net.torch.extractor import GraphExtractor as BuiltinGraphExtractor
 import graph_net.imp_util as imp_util
 
+NODE_COUNT_FILE = "./output/node_count.txt"
+
+def get_graph_node_count(gm: torch.fx.GraphModule):
+    """返回任意 GraphModule 的节点数量（包括子图）"""
+    return len(list(gm.graph.nodes))
+
+def write_global_node_count(n: int):
+    os.makedirs(os.path.dirname(NODE_COUNT_FILE), exist_ok=True)
+    with open(NODE_COUNT_FILE, "w") as f:
+        f.write(str(n))
+
+def write_subgraph_node_count(path: str, n: int):
+    """写入 path/node_count.txt"""
+    os.makedirs(path, exist_ok=True)
+    file = os.path.join(path, "node_count.txt")
+    with open(file, "w") as f:
+        f.write(str(n))
 
 class GraphExtractor:
     def __init__(
@@ -47,6 +64,10 @@ def make_config(
         }
 
     def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
+        node_count = len(list(gm.graph.nodes))
+        write_global_node_count(node_count)
+        print(f"[naive_graph_decomposer] 总节点数: {node_count} -> {NODE_COUNT_FILE}")
+
         config = {
             k: v
             for k, v in self.config.items()
@@ -83,6 +104,17 @@ def __init__(self, parent_graph_extractor, submodule, seq_no):
     def forward(self, *args):
         if not self.extracted:
             if self.need_extract(self.submodule, args):
+                out_dir = os.path.join(
+                    self.parent_graph_extractor.config["output_dir"],
+                    f"{self.parent_graph_extractor.name}_{self.seq_no}"
+                )
+
+                # 子图节点数
+                node_count = len(list(self.submodule.graph.nodes))
+                write_subgraph_node_count(out_dir, node_count)
+
+                print(f"[naive_graph_decomposer] 子图 {self.seq_no} 节点数 = {node_count} -> {out_dir}/node_count.txt")
+                
                 self.builtin_extractor(self.submodule, args)
             self.extracted = True
         return self.submodule(*args)