fix style

Author: TelGome

graph_net/test/graph_variable_rename_test.sh

@@ -29,6 +29,7 @@ python3 -m graph_net.model_path_handler --model-path samples/$MODEL_PATH_IN_SAMP
 
 test_compiler_config_json_str=$(cat <<EOF
 {
+    "model_path_prefix": "$GRAPH_NET_ROOT",
     "renamed_root": "$WORKSPACE"
 }
 EOF

graph_net/torch/backend/graph_variable_renamer_validator_backend.py

@@ -1,8 +1,9 @@
 import torch
 from pathlib import Path
-from typing import Dict, Any
-from graph_net.imp_util import load_module
+from typing import Dict
 from graph_net.tensor_meta import TensorMeta
+import os
+import importlib.util
 
 
 class RenamedModelAdapter(torch.nn.Module):
@@ -14,104 +15,75 @@ def __init__(self, renamed_model: torch.nn.Module, mapping: Dict[str, str]):
             self.__graph_net_file_path__ = renamed_model.__graph_net_file_path__
 
     def forward(self, **kwargs):
+        new_kwargs = self._convert_by_name_mapping(kwargs)
+        return self.model(**new_kwargs)
+
+    def _convert_by_name_mapping(self, kwargs):
         new_kwargs = {}
         for old_name, value in kwargs.items():
             if old_name in self.mapping:
                 new_name = self.mapping[old_name]
                 new_kwargs[new_name] = value
-        return self.model(**new_kwargs)
+        return new_kwargs
 
 
 class GraphVariableRenamerValidatorBackend:
-    def __init__(self, config: Dict[str, Any] = None):
-        self.config = config or {}
-
-    def _get_rename_mapping(self, model_dir: Path) -> Dict[str, str]:
+    def _get_rename_mapping(self, model_dir: Path):
         mapping = {}
-        if not model_dir.exists():
-            print(f"[ValidatorBackend] Error: Model dir does not exist: {model_dir}")
-            return mapping
-
         for meta_file in ["input_meta.py", "weight_meta.py"]:
             meta_path = model_dir / meta_file
-            if meta_path.exists():
-                try:
-                    metas = TensorMeta.unserialize_from_py_file(str(meta_path))
-                    for m in metas:
-                        if m.original_name:
-                            mapping[m.original_name] = m.name
-                except Exception as e:
-                    print(
-                        f"[ValidatorBackend] Warning: Failed to parse {meta_path}: {e}"
-                    )
+            if not meta_path.exists():
+                continue
+            metas = TensorMeta.unserialize_from_py_file(str(meta_path))
+            for m in metas:
+                if m.original_name:
+                    mapping[m.original_name] = m.name
         return mapping
 
-    def _load_renamed_model(
-        self, model_dir: Path, device: torch.device
-    ) -> torch.nn.Module:
-        model_py_path = model_dir / "model.py"
-        if not model_py_path.exists():
-            raise FileNotFoundError(f"Renamed model not found at {model_py_path}")
-
-        py_module = load_module(str(model_py_path))
-
-        if not hasattr(py_module, "GraphModule"):
-            raise ValueError(f"GraphModule class not found in {model_py_path}")
-
-        GraphModule = getattr(py_module, "GraphModule")
-        GraphModule.__graph_net_file_path__ = str(model_py_path)
-
-        model = GraphModule()
-        model.to(device)
-        model.eval()
-        return model
-
-    def __call__(self, original_model: torch.nn.Module) -> torch.nn.Module:
-        renamed_root = self.config.get("renamed_root")
-        if not renamed_root:
-            raise ValueError("Config 'renamed_root' is missing!")
-
-        default_prefix = str(Path(__file__).resolve().parent.parent.parent.parent)
-        model_path_prefix = self.config.get("model_path_prefix", default_prefix)
-
-        if not hasattr(original_model, "__graph_net_file_path__"):
-            raise ValueError("Original model missing __graph_net_file_path__")
-
-        orig_abs_path = Path(original_model.__class__.__graph_net_file_path__).resolve()
-        orig_model_dir = orig_abs_path.parent
-
-        try:
-            rel_model_path = orig_model_dir.relative_to(model_path_prefix)
-        except ValueError:
-            print(
-                f"[ValidatorBackend] Warning: Model path {orig_model_dir} is not under prefix {model_path_prefix}. Fallback to leaf name."
-            )
-            rel_model_path = orig_model_dir.name
-
-        renamed_model_dir = Path(renamed_root) / rel_model_path
-
-        print(f"[ValidatorBackend] Original Path: {orig_model_dir}")
-        print(f"[ValidatorBackend] Relative Path: {rel_model_path}")
-        print(f"[ValidatorBackend] Loading Renamed: {renamed_model_dir}")
-
-        try:
-            device = next(original_model.parameters()).device
-        except StopIteration:
-            device = torch.device("cpu")
-
-        renamed_core_model = self._load_renamed_model(renamed_model_dir, device)
-        mapping = self._get_rename_mapping(renamed_model_dir)
-
-        if not mapping:
-            print(
-                f"[ValidatorBackend] Warning: Mapping is empty for {rel_model_path}. Check input_meta.py generation."
-            )
-
-        adapter = RenamedModelAdapter(renamed_core_model, mapping)
-        adapter.to(device)
-        adapter.eval()
-
-        return adapter
+    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 _make_config(
+        self,
+        model_path_prefix: str,
+        renamed_root: str,
+        renamed_dentry: str = "_renamed",
+    ):
+        return {
+            "model_path_prefix": model_path_prefix,
+            "renamed_root": renamed_root,
+            "renamed_dentry": renamed_dentry,
+        }
+
+    def __call__(self, model: torch.nn.Module) -> torch.nn.Module:
+        config = self._make_config(**self.config)
+        model_path = os.path.dirname(model.__class__.__graph_net_file_path__)
+        model_name = os.path.basename(model_path)
+        renamed_dir_name = f"{model_name}_renamed"
+        renamed_model_dir = os.path.join(config["renamed_root"], renamed_dir_name)
+
+        print(f"[GraphVariableRenamerValidatorBackend] Processing: {model_name}")
+        print(
+            f"[GraphVariableRenamerValidatorBackend] Loading from: {renamed_model_dir}"
+        )
+
+        device = model.__class__.__graph_net_device__
+        renamed_model = self._load_model_instance(renamed_model_dir, device)
+        mapping = self._get_rename_mapping(Path(renamed_model_dir))
+        assert (
+            mapping
+        ), f"Mapping is empty for {renamed_dir_name} at {renamed_model_dir}"
+        adapter = RenamedModelAdapter(renamed_model, mapping)
+        return adapter.eval()
 
     def synchronize(self):
         if torch.cuda.is_available():

graph_net/torch/graph_variable_renamer.py

@@ -79,8 +79,10 @@ def __call__(self, rel_model_path):
         module, inputs = get_torch_module_and_inputs(src_model_path)
         gm = parse_sole_graph_module(module, inputs)
         gm = self.rename_graph_variables(gm, inputs, src_model_path)
+        model_name = os.path.basename(rel_model_path.rstrip(os.sep))
+        new_rel_path = f"{model_name}_renamed"
         dst_model_path = os.path.realpath(
-            os.path.join(self.config["output_dir"], rel_model_path)
+            os.path.join(self.config["output_dir"], new_rel_path)
         )
         Path(dst_model_path).parent.mkdir(parents=True, exist_ok=True)
         shutil.copytree(src_model_path, dst_model_path, dirs_exist_ok=True)
@@ -158,55 +160,47 @@ def _get_input_names_from_signature(self, module):
     def rename_graph_variables(
         self, gm: torch.fx.GraphModule, sample_inputs, model_path
     ):
-        in_cnt = 0
-        w_cnt = 0
-        tmp_cnt = 0
-
+        counters = {"in": 0, "w": 0, "tmp": 0}
+        # graph may not have input, only contain weights
         arg_iter = iter(sample_inputs) if sample_inputs else iter([])
-
         for node in gm.graph.nodes:
-            old_name = node.name
-
-            if "original_name" not in node.meta:
-                node.meta["original_name"] = node.name
-
-            new_name = old_name
-            should_update_target = False
-
-            if node.op == "placeholder":
-                real_arg = next(arg_iter)
-                is_weight = not self.data_input_predicator(model_path, node.name)
-                if node.type is not None:
-                    if isinstance(node.type, type) and issubclass(
-                        node.type, torch.nn.parameter.Parameter
-                    ):
-                        is_weight = True
-                elif real_arg is not None:
-                    if isinstance(real_arg, torch.nn.Parameter):
-                        is_weight = True
-
-                if is_weight:
-                    new_name = f"w_{w_cnt}"
-                    w_cnt += 1
-                else:
-                    new_name = f"in_{in_cnt}"
-                    in_cnt += 1
-
-                should_update_target = True
-
-            elif node.op == "get_attr":
-                node.name = f"w_{w_cnt}"
-                w_cnt += 1
-
-            elif node.op != "output":
-                node.name = f"tmp_{tmp_cnt}"
-                tmp_cnt += 1
-
-            if new_name != old_name:
-                node.name = new_name
-                if should_update_target:
-                    node.target = new_name
-
+            self._process_single_node(node, arg_iter, counters, model_path)
         gm.graph.lint()
         gm.recompile()
         return gm
+
+    def _process_single_node(self, node, arg_iter, counters, model_path):
+        if "original_name" not in node.meta:
+            node.meta["original_name"] = node.name
+        if node.op == "placeholder":
+            self._handle_placeholder(node, arg_iter, counters, model_path)
+        elif node.op == "get_attr":
+            self._apply_rename(node, "w", counters)
+        elif node.op != "output":
+            self._apply_rename(node, "tmp", counters)
+
+    def _handle_placeholder(self, node, arg_iter, counters, model_path):
+        real_arg = next(arg_iter, None)
+        is_weight = self._is_weight_node(node, real_arg, model_path)
+        prefix = "w" if is_weight else "in"
+        self._apply_rename(node, prefix, counters, update_target=True)
+
+    def _apply_rename(self, node, prefix, counters, update_target=False):
+        new_name = f"{prefix}_{counters[prefix]}"
+        counters[prefix] += 1
+        if node.name != new_name:
+            node.name = new_name
+            if update_target:
+                node.target = new_name
+
+    def _is_weight_node(self, node, real_arg, model_path):
+        if not self.data_input_predicator(model_path, node.name):
+            return True
+        if node.type is not None:
+            if isinstance(node.type, type) and issubclass(
+                node.type, torch.nn.parameter.Parameter
+            ):
+                return True
+        if real_arg is not None and isinstance(real_arg, torch.nn.Parameter):
+            return True
+        return False