Naive graph decomposer (#338)

* support checking model redundancy

* revert change of vision_model_test

* reformat python code.

* reformat bert_model_test.py and utils.py

* minor fix

* fix failed check by comparing directories after os.path.realpath()

* fix bugs in check_validate.sh

* set dynamic=False in single_device_runner.py

* reset graph hash

* naive_graph_decomposer

Author: lixinqi

graph_net/test/naive_graph_decomposer_test.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+# input model path
+MODEL_PATH_IN_SAMPLES=/timm/resnet18 
+# output model path
+OUTPUT_DIR=/tmp/naive_decompose_workspace
+
+mkdir -p $OUTPUT_DIR
+# extract subgraph 0-8, 8-16
+export GRAPH_NET_NAIVE_DECOMPOSER_SPLIT_POS=0,8,16
+export GRAPH_NET_EXTRACT_WORKSPACE=$OUTPUT_DIR 
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+python3 -m graph_net.torch.single_device_runner --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --enable-extract True --extract-name resnet18 --dump-graph-hash-key --custom-extractor-path=$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py

graph_net/torch/decompose_util.py

@@ -33,17 +33,16 @@ def convert_to_submodules_graph(
     split_positions = [
         max(0, min(pos, len(submodules_body_nodes))) for pos in split_positions
     ]
-    submodule_ranges = [
-        (start, end)
+    range_idx2submodule_body_nodes = [
+        submodules_body_nodes[start:end]
         for i in range(len(split_positions) - 1)
         for start in [split_positions[i]]
         for end in [split_positions[i + 1]]
         if end > start
     ]
 
     def get_body_nodes(range_idx):
-        start, end = submodule_ranges[range_idx]
-        return submodules_body_nodes[start:end]
+        return range_idx2submodule_body_nodes[range_idx]
 
     def get_name2sub_submodule():
         used_module_names = set(
@@ -55,15 +54,28 @@ def get_name2sub_submodule():
             if name in used_module_names
         }
 
-    for range_idx in range(len(submodule_ranges)):
-        start_node_idx, end_node_idx = submodule_ranges[range_idx]
+    def get_start_node_idx(range_idx):
+        start_node = get_body_nodes(range_idx)[0]
+        for i, node in enumerate(original_gm.graph.nodes):
+            if node == start_node:
+                return i
+        raise NotImplementedError("Dead code.")
+
+    def get_end_node_idx(range_idx):
+        last_node = get_body_nodes(range_idx)[-1]
+        for i, node in enumerate(original_gm.graph.nodes):
+            if node == last_node:
+                return i + 1
+        raise NotImplementedError("Dead code.")
+
+    for range_idx in range(len(range_idx2submodule_body_nodes)):
         (
             submodule_input_nodes,
             submodule_output_nodes,
         ) = _get_submodule_inputs_and_outputs(
             original_gm=original_gm,
-            start_node_idx=(num_placeholders + start_node_idx),
-            end_node_idx=(num_placeholders + end_node_idx),
+            start_node_idx=get_start_node_idx(range_idx),
+            end_node_idx=get_end_node_idx(range_idx),
         )
 
         def get_input_nodes(range_idx):

graph_net/torch/extractor.py

@@ -125,7 +125,13 @@ def try_rename_placeholder(node):
         return gm.forward
 
 
-def extract(name, dynamic=True, mut_graph_codes=None, placeholder_auto_rename=False):
+def extract(
+    name,
+    dynamic=True,
+    mut_graph_codes=None,
+    placeholder_auto_rename=False,
+    custom_extractor_path=None,
+):
     """
     Extract computation graphs from PyTorch nn.Module.
     The extracted computation graphs will be saved into directory of env var $GRAPH_NET_EXTRACT_WORKSPACE.
@@ -194,9 +200,19 @@ def forward(self, s0 : torch.SymInt, L_x_ : torch.Tensor):
         >>>
     """
 
+    def get_graph_extractor_cls():
+        if custom_extractor_path is None:
+            return GraphExtractor
+        import importlib.util as imp
+
+        spec = imp.spec_from_file_location("graph_extractor", custom_extractor_path)
+        graph_extractor = imp.module_from_spec(spec)
+        spec.loader.exec_module(graph_extractor)
+        return graph_extractor.GraphExtractor
+
     def wrapper(model: torch.nn.Module):
         assert isinstance(model, torch.nn.Module), f"{type(model)=}"
-        extractor = GraphExtractor(
+        extractor = get_graph_extractor_cls()(
             name, dynamic, mut_graph_codes, placeholder_auto_rename
         )
         # return torch.compile(backend=extractor, dynamic=dynamic)

graph_net/torch/naive_graph_decomposer.py

@@ -0,0 +1,63 @@
+import os
+import torch
+import json
+import shutil
+from typing import Union, Callable
+from graph_net.torch import utils
+from graph_net.torch.decompose_util import convert_to_submodules_graph
+from graph_net.torch.extractor import GraphExtractor as BuiltinGraphExtractor
+
+
+class GraphExtractor:
+    def __init__(
+        self, name, dynamic, mut_graph_codes=None, placeholder_auto_rename=False
+    ):
+        self.subgraph_counter = 0
+        self.name = name
+        self.dynamic = dynamic
+        self.mut_graph_codes = mut_graph_codes
+        self.placeholder_auto_rename = placeholder_auto_rename
+        self.workspace_path = os.environ.get("GRAPH_NET_EXTRACT_WORKSPACE")
+        if not self.workspace_path:
+            raise EnvironmentError(
+                "Environment variable 'GRAPH_NET_EXTRACT_WORKSPACE' is not set."
+            )
+        split_pos_str = os.environ.get("GRAPH_NET_NAIVE_DECOMPOSER_SPLIT_POS")
+        if split_pos_str is None:
+            raise EnvironmentError(
+                "Environment variable 'GRAPH_NET_NAIVE_DECOMPOSER_SPLIT_POS' is not set."
+            )
+        self.split_positions = [int(pos) for pos in split_pos_str.split(",")]
+
+    def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
+        return convert_to_submodules_graph(
+            gm,
+            split_positions=self.split_positions,
+            submodule_hook=self.get_naive_decomposer_extractor,
+            group_head_and_tail=False,
+        )
+
+    def get_naive_decomposer_extractor(self, submodule, seq_no):
+        return NaiveDecomposerExtractor(self, submodule, seq_no)
+
+
+class NaiveDecomposerExtractor(torch.nn.Module):
+    def __init__(self, parent_graph_extractor, submodule, seq_no):
+        super().__init__()
+        self.parent_graph_extractor = parent_graph_extractor
+        self.submodule = submodule
+        self.seq_no = seq_no
+        self.extracted = False
+        name = f"{parent_graph_extractor.name}_{self.seq_no}"
+        self.builtin_extractor = BuiltinGraphExtractor(
+            name=name,
+            dynamic=False,
+            mut_graph_codes=[],
+            placeholder_auto_rename=parent_graph_extractor.placeholder_auto_rename,
+        )
+
+    def forward(self, *args):
+        if not self.extracted:
+            self.builtin_extractor(self.submodule, args)
+            self.extracted = True
+        return self.submodule(*args)

graph_net/torch/single_device_runner.py

@@ -59,7 +59,13 @@ def main(args):
         print(f"{model_path=}")
         if args.enable_extract:
             assert args.extract_name is not None
-            kwargs = dict(name=args.extract_name, dynamic=False, **dump_graph_options)
+
+            kwargs = dict(
+                name=args.extract_name,
+                dynamic=False,
+                custom_extractor_path=args.custom_extractor_path,
+                **dump_graph_options,
+            )
             model = extract(**kwargs)(model)
 
         inputs_params = utils.load_converted_from_text(f"{model_path}")
@@ -110,5 +116,12 @@ def main(args):
         default=None,
         help="Extracted graph's name",
     )
+    parser.add_argument(
+        "--custom-extractor-path",
+        type=str,
+        required=False,
+        default=None,
+        help="Custom extractor python file path",
+    )
     args = parser.parse_args()
     main(args=args)