modify script, rename files and variables

Author: roll-away

graph_net/test/naive_decomposer_and_post_extract_process_test.sh

@@ -1,10 +1,5 @@
 #!/bin/bash
-SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
-GRAPH_NET_DIR=$(dirname "$SCRIPT_DIR")
-PROJECT_ROOT=$(dirname "$GRAPH_NET_DIR")
-
-# 将项目根目录加入Python路径
-export PYTHONPATH="$PROJECT_ROOT:$PYTHONPATH"
+# bash graph_net/test/naive_decomposer_and_post_extract_process_test.sh
 
 GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
 os.path.dirname(graph_net.__file__))")
@@ -19,12 +14,13 @@ decorator_config_json_str=$(cat <<EOF
         "name": "$MODEL_NAME",
         "custom_extractor_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
         "custom_extractor_config": {
-            "output_dir": "/work/.BCloud/countkernels/",
+            "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",
             "filter_config": {},
-            "post_extract_process_path":"$GRAPH_NET_ROOT/torch/post_extract_process.py"
+            "post_extract_process_path":"$GRAPH_NET_ROOT/torch/post_extract_process_count_kernels.py",
+            "post_extract_process_class_name": "PostExtractProcess"
         }
     }
 }

graph_net/torch/naive_graph_decomposer.py

@@ -33,7 +33,7 @@ def make_config(
         filter_path=None,
         filter_config=None,
         post_extract_process_path=None,
-        post_extract_process_config=None,
+        post_extract_process_class_name=None,
     ):
         for pos in split_positions:
             assert isinstance(
@@ -47,7 +47,7 @@ def make_config(
             "filter_path": filter_path,
             "filter_config": filter_config if filter_config is not None else {},
             "post_extract_process_path": post_extract_process_path,
-            "post_extract_process_config": post_extract_process_config,
+            "post_extract_process_class_name": post_extract_process_class_name,
         }
 
     def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
@@ -75,7 +75,7 @@ def __init__(self, parent_graph_extractor, submodule, seq_no):
         self.seq_no = seq_no
         self.extracted = False
         name = f"{parent_graph_extractor.name}_{self.seq_no}"
-        self.modelname = name
+        self.model_name = name
         self.builtin_extractor = BuiltinGraphExtractor(
             name=name,
             dynamic=False,
@@ -103,7 +103,7 @@ def need_extract(self, gm, sample_inputs):
 
     def _post_extract_process(self):
         model_path = os.path.join(
-            self.parent_graph_extractor.config["output_dir"], self.modelname
+            self.parent_graph_extractor.config["output_dir"], self.model_name
         )
         return self.post_extract_process(model_path)
 
@@ -117,4 +117,4 @@ def make_post_extract_process(self, config):
         if config["post_extract_process_path"] is None:
             return None
         module = imp_util.load_module(config["post_extract_process_path"])
-        return module.PostExtractProcess(config["post_extract_process_config"])
+        return module.PostExtractProcess(config["post_extract_process_path"])

graph_net/torch/post_extract_process.py …ch/post_extract_process_count_kernels.pygraph_net/torch/post_extract_process.py renamed to graph_net/torch/post_extract_process_count_kernels.py graph_net/torch/post_extract_process.py renamed to graph_net/torch/post_extract_process_count_kernels.py

@@ -27,7 +27,7 @@ def __call__(self, model_path=None):
 
         model(**state_dict)
         compiled_model = torch.compile(model)
-        compiled_num_of_kernels = count_kernels(model, state_dict)
+        compiled_num_of_kernels = count_kernels(compiled_model, state_dict)
         if compiled_num_of_kernels == 1:
             print(model_path, "can be fully integrated")
             return True