fix

Author: fangfangssj

graph_net/test/chain_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/chain_naive_decompose_workspace",
         "split_positions": [8, 16, 32],
         "group_head_and_tail": true,
         "chain_style": true
@@ -23,4 +20,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 $MODEL_NAME --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 resnet18 --dump-graph-hash-key --extractor-config=$EXTRACTOR_CONFIG

graph_net/test/decomposer_validator_test.sh

@@ -1,25 +1,40 @@
 #!/bin/bash
 
-if [ -z "$GRAPH_NET_BENCHMARK_PATH" ]; then
-    GRAPH_NET_BENCHMARK_PATH="$(pwd)/graphnet_benchmark"
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(os.path.dirname(graph_net.__file__))")
+
+if [ -z "$GRAPH_NET_DECOMPOSE_PATH" ]; then
+    GRAPH_NET_DECOMPOSE_PATH="$(pwd)/graphnet_decompose"
 fi
 
-FILE_PATH=$GRAPH_NET_BENCHMARK_PATH/decomposer
+MODEL_PATH_IN_SAMPLES=/timm/resnet18
+MODEL_NAME=$(basename "$MODEL_PATH_IN_SAMPLES")
+OUTPUT_DIR="${GRAPH_NET_DECOMPOSE_PATH:-$(pwd)}"
+cp -r "$GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES" "$OUTPUT_DIR/"
+
+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",
+        "split_positions": [8, 16, 32],
+        "group_head_and_tail": true,
+        "chain_style": true
+    }
+}
+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
+
+FILE_PATH=$GRAPH_NET_DECOMPOSE_PATH/decomposer
 mkdir -p "$(dirname "$FILE_PATH/log.log")"
-
-MODEL_PATH="./todo_works/range_decomposer_validator/test/resnet18"
+MODEL_PATH="$GRAPH_NET_DECOMPOSE_PATH/$MODEL_NAME"
 
 python -m graph_net.torch.test_compiler \
     --model-path $MODEL_PATH \
     --compiler range_decomposer_validator \
     --device cuda > "$FILE_PATH/log.log" 2>&1
 
-if [ $? -ne 0 ]; then
-    echo "Error: decomposer_validator execution failed"
-    echo "Please check the log file: $FILE_PATH/log.log"
-    exit 1
-fi
-
 python -m graph_net.log2json \
     --log-file "$FILE_PATH/log.log" \
     --output-dir "$FILE_PATH/JSON_results/"

graph_net/torch/backend/range_decomposer_validator_backend.py

@@ -12,54 +12,10 @@ class ComposedModel(nn.Module):
     def __init__(self, graph: nn.Module, subgraph: List[nn.Module]):
         super().__init__()
         self.graph = graph
-        self.subgraph = nn.ModuleList(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):
-            return arg.name
-        if isinstance(arg, (list, tuple)):
-            return type(arg)(self._serialize_arg(elem) for elem in arg)
-        if isinstance(arg, dict):
-            return {
-                self._serialize_arg(k): self._serialize_arg(v) for k, v in arg.items()
-            }
-        return arg
-
-    def _extract_operators_from_graph(
-        self, gm: nn.Module, example_inputs: List[torch.Tensor] = None
-    ) -> List[Dict[str, Any]]:
-        operator_list = []
-        for node in gm.graph.nodes:
-            if node.op in ("call_method", "call_function", "call_module"):
-                operator_info = {
-                    "op_type": node.op,
-                    "target": node.target,
-                    "kwargs": self._serialize_arg(node.kwargs),
-                }
-
-                if isinstance(node.target, Callable):
-                    try:
-                        operator_info["target_name"] = node.target.__name__
-                    except AttributeError:
-                        operator_info["target_name"] = str(node.target)
-                else:
-                    operator_info["target_name"] = str(node.target)
-
-                operator_list.append(operator_info)
-
-        return operator_list
-
-    def extract_compiler(self, gm: torch.fx.GraphModule, inputs: List[torch.Tensor]):
-        operator = self._extract_operators_from_graph(gm, inputs)
-        self.extract_node.append(operator)
-        return gm.forward
+        self.subgraphs = nn.ModuleList(subgraph)
 
     def forward(self, **kwargs):
-        self.graph_model(**kwargs)
-        graph_node_list = list(itertools.chain.from_iterable(self.extract_node))
-        self.extract_node = []
+        self.graph(**kwargs)
 
         subgraph_intput = {
             key.replace("L", "l_l", 1): value
@@ -68,31 +24,11 @@ def forward(self, **kwargs):
         }
 
         output = None
-        for subgraph_model in self.subgraph:
-            compiled_model = torch.compile(
-                subgraph_model, backend=self.extract_compiler
-            )
-
+        for subgraph in self.subgraphs:
             if output is None:
-                output = compiled_model(**subgraph_intput)
+                output = subgraph(**subgraph_intput)
             else:
-                output = compiled_model(*output)
-
-        subgraph_node_list = list(itertools.chain.from_iterable(self.extract_node))
-        self.extract_node = []
-
-        if graph_node_list != subgraph_node_list:
-            diff_in_graph = [
-                item for item in graph_node_list if item not in subgraph_node_list
-            ]
-            diff_in_subgraph = [
-                item for item in subgraph_node_list if item not in graph_node_list
-            ]
-
-            error_msg = f"Subgraph segmentation verification failed\n"
-            error_msg += f"Nodes in graph but not in subgraph: {diff_in_graph}\n"
-            error_msg += f"Nodes in subgraph but not in graph: {diff_in_subgraph}"
-            raise ValueError(error_msg)
+                output = subgraph(*output)
 
         return output