fix

Author: fangfangssj

graph_net/test/chain_naive_graph_decomposer_test.sh

@@ -4,12 +4,15 @@ GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
 os.path.dirname(graph_net.__file__))")
 
 # input model path
-MODEL_PATH_IN_SAMPLES=/timm/resnet18 
+MODEL_PATH_IN_SAMPLES=/timm/resnet18
+MODEL_NAME=$(basename "$MODEL_PATH_IN_SAMPLES")
+OUTPUT_DIR="${NAIVE_DECOMPOSE_WORKSPACE:-$(pwd)/naive_decompose_workspace}"
+
 extractor_config_json_str=$(cat <<EOF
 {
     "custom_extractor_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
     "custom_extractor_config": {
-        "output_dir": "/tmp/chain_naive_decompose_workspace",
+        "output_dir": "$OUTPUT_DIR/${MODEL_NAME}_decomposed",
         "split_positions": [8, 16, 32],
         "group_head_and_tail": true,
         "chain_style": true
@@ -20,4 +23,4 @@ EOF
 EXTRACTOR_CONFIG=$(echo $extractor_config_json_str | base64 -w 0)
 
 mkdir -p /tmp/naive_decompose_workspace
-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 --extractor-config=$EXTRACTOR_CONFIG
+python3 -m graph_net.torch.single_device_runner --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --enable-extract True --extract-name $MODEL_NAME --dump-graph-hash-key --extractor-config=$EXTRACTOR_CONFIG

graph_net/test/decomposer_validator_test.sh

@@ -1,13 +1,13 @@
 #!/bin/bash
 
 if [ -z "$GRAPH_NET_BENCHMARK_PATH" ]; then
-    GRAPH_NET_BENCHMARK_PATH="$(pwd)"
+    GRAPH_NET_BENCHMARK_PATH="$(pwd)/graphnet_benchmark"
 fi
 
 FILE_PATH=$GRAPH_NET_BENCHMARK_PATH/decomposer
 mkdir -p "$(dirname "$FILE_PATH/log.log")"
 
-MODEL_PATH="./todo_works/range_decomposer_validator/test/simple_CNN"
+MODEL_PATH="./todo_works/range_decomposer_validator/test/resnet18"
 
 python -m graph_net.torch.test_compiler \
     --model-path $MODEL_PATH \

graph_net/test/naive_graph_decomposer_test.sh

@@ -4,15 +4,12 @@ GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
 os.path.dirname(graph_net.__file__))")
 
 # input model path
-MODEL_PATH_IN_SAMPLES=/timm/resnet18
-MODEL_NAME=$(basename "$MODEL_PATH_IN_SAMPLES")
-OUTPUT_DIR="${NAIVE_DECOMPOSE_WORKSPACE:-$(pwd)/naive_decompose_workspace}"
-
+MODEL_PATH_IN_SAMPLES=/timm/resnet18 
 extractor_config_json_str=$(cat <<EOF
 {
     "custom_extractor_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
     "custom_extractor_config": {
-        "output_dir": "$OUTPUT_DIR/${MODEL_NAME}_decomposed",
+        "output_dir": "/tmp/naive_decompose_workspace",
         "split_positions": [8, 16, 32],
         "group_head_and_tail": true,
         "filter_path":"$GRAPH_NET_ROOT/torch/naive_subgraph_filter.py",
@@ -23,4 +20,5 @@ EOF
 )
 EXTRACTOR_CONFIG=$(echo $extractor_config_json_str | base64 -w 0)
 
-python3 -m graph_net.torch.single_device_runner --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --enable-extract True --extract-name $MODEL_NAME --dump-graph-hash-key --extractor-config=$EXTRACTOR_CONFIG
+mkdir -p /tmp/naive_decompose_workspace
+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 --extractor-config=$EXTRACTOR_CONFIG

…_validator/range_decomposer_validator.py …nd/range_decomposer_validator_backend.pytodo_works/range_decomposer_validator/range_decomposer_validator.py renamed to graph_net/torch/backend/range_decomposer_validator_backend.py todo_works/range_decomposer_validator/range_decomposer_validator.py renamed to graph_net/torch/backend/range_decomposer_validator_backend.py

@@ -13,11 +13,8 @@ def __init__(self, graph: nn.Module, subgraph: List[nn.Module]):
         super().__init__()
         self.graph = graph
         self.subgraph = nn.ModuleList(subgraph)
-        self.subgraph_param_names = [
-            list(inspect.signature(sm.forward).parameters.keys())
-            for sm in self.subgraph
-        ]
         self.extract_node = []
+        self.graph_model = torch.compile(self.graph, backend=self.extract_compiler)
 
     def _serialize_arg(self, arg: Any) -> Any:
         if isinstance(arg, torch.fx.Node):
@@ -39,7 +36,6 @@ def _extract_operators_from_graph(
                 operator_info = {
                     "op_type": node.op,
                     "target": node.target,
-                    "name": node.name,
                     "kwargs": self._serialize_arg(node.kwargs),
                 }
 
@@ -61,30 +57,26 @@ def extract_compiler(self, gm: torch.fx.GraphModule, inputs: List[torch.Tensor])
         return gm.forward
 
     def forward(self, **kwargs):
-        current_args = kwargs
-        compiled_model = torch.compile(self.graph, backend=self.extract_compiler)
-        compiled_model(**current_args)
+        self.graph_model(**kwargs)
         graph_node_list = list(itertools.chain.from_iterable(self.extract_node))
         self.extract_node = []
 
-        for i, (sm, param_names) in enumerate(
-            zip(self.subgraph, self.subgraph_param_names)
-        ):
-            call_kwargs = {}
-            if i > 0:
-                first_param_name = param_names[0]
-                call_kwargs[first_param_name] = current_args
-                remaining_params = param_names[1:]
-            else:
-                remaining_params = param_names
+        subgraph_intput = {
+            key.replace("L", "l_l", 1): value
+            for key, value in kwargs.items()
+            if key.startswith("L")
+        }
 
-            for name in remaining_params:
-                if name in kwargs:
-                    call_kwargs[name] = kwargs[name]
+        output = None
+        for subgraph_model in self.subgraph:
+            compiled_model = torch.compile(
+                subgraph_model, backend=self.extract_compiler
+            )
 
-            compiled_model = torch.compile(sm, backend=self.extract_compiler)
-            outputs = compiled_model(**call_kwargs)
-            current_args = outputs[0]
+            if output is None:
+                output = compiled_model(**subgraph_intput)
+            else:
+                output = compiled_model(*output)
 
         subgraph_node_list = list(itertools.chain.from_iterable(self.extract_node))
         self.extract_node = []
@@ -102,7 +94,7 @@ def forward(self, **kwargs):
             error_msg += f"Nodes in subgraph but not in graph: {diff_in_subgraph}"
             raise ValueError(error_msg)
 
-        return (current_args,)
+        return output
 
 
 class RangeDecomposerValidatorBackend:
@@ -119,7 +111,7 @@ def _load_model_instance(self, path: str, device: str) -> torch.nn.Module:
         return instance
 
     def __call__(self, model: torch.nn.Module) -> torch.nn.Module:
-        model_file_path = model.__class__.__file_path__
+        model_file_path = model.__class__.__graph_net_file_path__
         model_dir = os.path.dirname(model_file_path)
         decomposed_parent_dir = model_dir + "_decomposed"
         subgraph_paths = []
@@ -132,7 +124,7 @@ def __call__(self, model: torch.nn.Module) -> torch.nn.Module:
             f"[RangeDecomposerValidatorBackend] Found subgraphs: {[os.path.basename(p) for p in subgraph_paths]}"
         )
 
-        device = model.__class__.__device__
+        device = model.__class__.__graph_net_device__
         graph_instances = self._load_model_instance(model_dir, device)
         subgraph_instances = []
 

graph_net/torch/test_compiler.py

@@ -23,7 +23,7 @@
 from graph_net.torch.backend.blade_disc_backend import BladeDISCBackend
 from graph_net.torch.backend.nope_backend import NopeBackend
 from graph_net.torch.backend.unstable_to_stable_backend import UnstableToStableBackend
-from todo_works.range_decomposer_validator.range_decomposer_validator import (
+from graph_net.torch.backend.range_decomposer_validator_backend import (
     RangeDecomposerValidatorBackend,
 )
 from graph_net.test_compiler_util import generate_allclose_configs
@@ -69,8 +69,8 @@ def load_class_from_file(
     exec(compiled_code, module.__dict__)
 
     model_class = getattr(module, class_name, None)
-    setattr(model_class, "__file_path__", file_path)
-    setattr(model_class, "__device__", device)
+    setattr(model_class, "__graph_net_file_path__", file_path)
+    setattr(model_class, "__graph_net_device__", device)
     return model_class
 
 

todo_works/range_decomposer_validator/test/resnet18/graph_hash.txt

@@ -0,0 +1 @@
+248d46ebcf5bc02d3e72953ea430b5e18175b0419dbdbcd2479202497f58319d

todo_works/range_decomposer_validator/test/resnet18/graph_net.json

@@ -0,0 +1,7 @@
+{
+    "framework": "torch",
+    "num_devices_required": 1,
+    "num_nodes_required": 1,
+    "source": "timm",
+    "heuristic_tag": "computer_vision"
+}