Merge branch 'develop' into one_hot119

Author: lixinqi

graph_net/test/chain_naive_graph_decomposer_test.sh

@@ -20,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 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 resnet18 --dump-graph-hash-key --extractor-config=$EXTRACTOR_CONFIG

graph_net/test/decomposer_validator_test.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+
+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
+
+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="$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
+
+python -m graph_net.log2json \
+    --log-file "$FILE_PATH/log.log" \
+    --output-dir "$FILE_PATH/JSON_results/"
+
+python -m graph_net.plot_ESt \
+    --benchmark-path "$FILE_PATH/JSON_results/" \
+    --output-dir "$FILE_PATH"
+
+echo "=================================================="
+echo "Results saved in: $FILE_PATH/ES_result.png"
+echo ""
+echo "IMPORTANT: Please verify if the curve in ES_result.png is a straight line"
+echo "If the curve is NOT a straight line, please check the log file: $FILE_PATH/log.log"
+echo "=================================================="

graph_net/test/naive_graph_decomposer_test.sh

@@ -21,4 +21,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 resnet18 --dump-graph-hash-key --extractor-config=$EXTRACTOR_CONFIG

graph_net/torch/backend/range_decomposer_validator_backend.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+import os
+import sys
+import inspect
+import importlib.util
+import itertools
+from typing import List, Tuple, Dict, Any, Callable
+
+
+class ComposedModel(nn.Module):
+    def __init__(self, subgraph: List[nn.Module]):
+        super().__init__()
+        self.subgraphs = nn.ModuleList(subgraph)
+
+    def forward(self, **kwargs):
+        subgraph_intput = {
+            key.replace("L", "l_l", 1): value
+            for key, value in kwargs.items()
+            if key.startswith("L")
+        }
+
+        output = None
+        for subgraph in self.subgraphs:
+            if output is None:
+                output = subgraph(**subgraph_intput)
+            else:
+                output = subgraph(*output)
+
+        return output
+
+
+class RangeDecomposerValidatorBackend:
+    def _load_model_instance(self, path: str, device: str) -> torch.nn.Module:
+        class_name = "GraphModule"
+        model_file = os.path.join(path, "model.py")
+
+        spec = importlib.util.spec_from_file_location(class_name, model_file)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+
+        ModelClass = getattr(module, class_name)
+        instance = ModelClass().to(device)
+        return instance
+
+    def __call__(self, model: torch.nn.Module) -> torch.nn.Module:
+        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 = []
+        for name in sorted(os.listdir(decomposed_parent_dir)):
+            full_path = os.path.join(decomposed_parent_dir, name)
+            if os.path.isdir(full_path) and name[-1].isdigit():
+                subgraph_paths.append(full_path)
+
+        print(
+            f"[RangeDecomposerValidatorBackend] Found subgraphs: {[os.path.basename(p) for p in subgraph_paths]}"
+        )
+
+        device = model.__class__.__graph_net_device__
+        subgraph_instances = []
+
+        for path in subgraph_paths:
+            instance = self._load_model_instance(path, device)
+            subgraph_instances.append(instance)
+            dir_name = os.path.basename(path)
+            print(
+                f"[RangeDecomposerValidatorBackend] Loaded and instantiated '{dir_name}'"
+            )
+
+        composed_model = ComposedModel(subgraph_instances)
+        return composed_model.eval()
+
+    def synchronize(self):
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -204,7 +204,28 @@ def _impl_unstable_to_stable_one_hot(self, gm):
 
         return gm
 
-    # replace this line with modification code for task 121 (torch._C._set_grad_enabled)
+    def _impl_unstable_to_stable_set_grad_enabled(self, gm):
+        """
+        Convert torch._C._set_grad_enabled and torch._C.set_grad_enabled to torch.set_grad_enabled
+        """
+
+        def replace_in_graph(graph_mod):
+            for node in graph_mod.graph.nodes:
+                if node.op == "call_function":
+                    if "set_grad_enabled" in str(node.target):
+                        node.target = torch.set_grad_enabled
+            graph_mod.recompile()
+
+        modules = [gm]
+        modules += [
+            m
+            for _, m in gm.named_modules()
+            if isinstance(m, torch.fx.GraphModule) and m is not gm
+        ]
+        for m in modules:
+            replace_in_graph(m)
+
+        return gm
 
     # replace this line with modification code for task 122 (torch._C._log_api_usage_once)
 

graph_net/torch/fx_graph_serialize_util.py

@@ -143,7 +143,8 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
         # replace this line with modification code for task 117 (torch._C._linalg.linalg_norm)
         (r"torch\._C\._nn\.softplus\(", "torch.nn.functional.softplus("),
         (r"torch\._C\._nn\.one_hot\(", "torch.nn.functional.one_hot("),
-        # replace this line with modification code for task 121 (torch._C._set_grad_enabled)
+        (r"torch\._C\._set_grad_enabled\(", "torch.set_grad_enabled("),
+        (r"torch\._C\.set_grad_enabled\(", "torch.set_grad_enabled("),
         # replace this line with modification code for task 122 (torch._C._log_api_usage_once)
         # replace this line with modification code for task 123 (torch._C._nn.pad)
         # replace this line with modification code for task 125 (torch._C._nn.gelu)

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,6 +69,8 @@ def load_class_from_file(
     exec(compiled_code, module.__dict__)
 
     model_class = getattr(module, class_name, None)
+    setattr(model_class, "__graph_net_file_path__", file_path)
+    setattr(model_class, "__graph_net_device__", device)
     return model_class