Support use original tensor_meta to to recover the re-extracted samples.

Author: Xreki

graph_net/imp_util.py

@@ -1,3 +1,4 @@
+import os
 import importlib.util as imp
 
 
@@ -6,5 +7,5 @@ def load_module(path, name="unnamed"):
     module = imp.module_from_spec(spec)
     module.__file__ = path
     spec.loader.exec_module(module)
-    module.__graph_net_file_path__ = path
+    module.__graph_net_file_path__ = os.path.normpath(path)
     return module

graph_net/paddle/graph_meta_restorer.py

@@ -0,0 +1,153 @@
+import os
+from graph_net import path_utils
+from graph_net.paddle import utils
+
+
+class GraphMetaRestorer:
+    def __init__(self, config, parent_model_path):
+        self.config = config
+        self.parent_model_path = parent_model_path
+        print(f"parent_model_path: {self.parent_model_path}")
+
+        assert path_utils.is_single_model_dir(
+            parent_model_path
+        ), f"{parent_model_path=} is not a graphnet sample."
+        (
+            parent_weight_meta_classes,
+            parent_input_meta_classes,
+        ) = self._load_weight_and_input_meta_classes(parent_model_path)
+        self.original_name2parent_weight_meta_class = self._convert_to_dict(
+            parent_weight_meta_classes
+        )
+        self.original_name2parent_input_meta_class = self._convert_to_dict(
+            parent_input_meta_classes
+        )
+
+    def __call__(self, model_path):
+        assert path_utils.is_single_model_dir(
+            model_path
+        ), f"{model_path=} is not a graphnet sample."
+        (
+            weight_meta_classes,
+            input_meta_classes,
+        ) = self._load_weight_and_input_meta_classes(model_path)
+
+        assert self.config["update_inplace"]
+        is_weight_meta_fully_updated = self._update_by_original_name(
+            weight_meta_classes, self.original_name2parent_weight_meta_class
+        )
+        if is_weight_meta_fully_updated:
+            new_weight_meta_codes = []
+            for meta_class in weight_meta_classes:
+                new_weight_meta_codes.append(
+                    self._generate_py_code_from_meta_class(meta_class)
+                )
+
+            weight_meta_file_path = os.path.join(model_path, "weight_meta.py")
+            if self.config["update_inplace"]:
+                print(f"[GraphMetaRestorer] Update {weight_meta_file_path}")
+                with open(weight_meta_file_path, "w") as f:
+                    f.write("\n\n".join(new_weight_meta_codes))
+
+        is_input_meta_fully_updated = self._update_by_tensor_spec(
+            input_meta_classes, self.original_name2parent_input_meta_class
+        )
+        if is_input_meta_fully_updated:
+            new_input_meta_codes = []
+            for meta_class in input_meta_classes:
+                new_input_meta_codes.append(
+                    self._generate_py_code_from_meta_class(meta_class)
+                )
+
+            input_meta_file_path = os.path.join(model_path, "input_meta.py")
+            if self.config["update_inplace"]:
+                print(f"[GraphMetaRestorer] Update {input_meta_file_path}")
+                with open(input_meta_file_path, "w") as f:
+                    f.write("\n\n".join(new_input_meta_codes))
+
+    def _load_weight_and_input_meta_classes(self, model_path):
+        weight_meta_file_path = os.path.join(model_path, "weight_meta.py")
+        weight_meta_classes = [
+            meta_class
+            for (name, meta_class) in utils.get_meta_classes(weight_meta_file_path)
+        ]
+
+        input_meta_file_path = os.path.join(model_path, "input_meta.py")
+        input_meta_classes = [
+            meta_class
+            for (name, meta_class) in utils.get_meta_classes(input_meta_file_path)
+        ]
+
+        return weight_meta_classes, input_meta_classes
+
+    def _convert_to_dict(self, meta_classes):
+        original_name2meta_class = {}
+        for meta_class in meta_classes:
+            assert meta_class.original_name not in original_name2meta_class.keys()
+            original_name2meta_class[meta_class.original_name] = meta_class
+        return original_name2meta_class
+
+    def _update_tensor_meta(self, meta_class, parent_meta_class):
+        if (
+            parent_meta_class
+            and meta_class.dtype == parent_meta_class.dtype
+            and meta_class.shape == parent_meta_class.shape
+        ):
+            for attr_name in ["max_val", "min_val", "mean", "std", "data"]:
+                if hasattr(meta_class, attr_name) or hasattr(
+                    parent_meta_class, attr_name
+                ):
+                    attr_value = getattr(parent_meta_class, attr_name, None)
+                    setattr(meta_class, attr_name, attr_value)
+            return True
+        return False
+
+    def _update_by_original_name(self, meta_classes, original_name2parent_meta_class):
+        updated_class_names = set()
+        for meta_class in meta_classes:
+            if not meta_class.original_name:
+                continue
+
+            parent_meta_class = original_name2parent_meta_class.get(
+                meta_class.original_name, None
+            )
+            if self._update_tensor_meta(meta_class, parent_meta_class):
+                updated_class_names.add(meta_class.name)
+
+        print(
+            f"[GraphMetaRestorer] {len(updated_class_names)}/{len(meta_classes)} classes are updated."
+        )
+        return len(meta_classes) == len(updated_class_names)
+
+    def _update_by_tensor_spec(self, meta_classes, original_name2parent_meta_class):
+        updated_class_names = set()
+        for meta_class in meta_classes:
+            matched_parent_meta_class = [
+                parent_meta_class
+                for parent_meta_class in original_name2parent_meta_class.values()
+                if meta_class.dtype == parent_meta_class.dtype
+                and meta_class.shape == parent_meta_class.shape
+            ]
+            if len(matched_parent_meta_class) == 1:
+                self._update_tensor_meta(meta_class, matched_parent_meta_class[0])
+                updated_class_names.add(meta_class.name)
+
+        print(
+            f"[GraphMetaRestorer] {len(updated_class_names)}/{len(meta_classes)} classes are updated."
+        )
+        return len(meta_classes) == len(updated_class_names)
+
+    def _generate_py_code_from_meta_class(self, meta_class):
+        lines = [f"class {meta_class.__name__}:"]
+        members = vars(meta_class)
+        members = {k: v for k, v in members.items() if not k.startswith("__")}
+
+        if not members:
+            return lines[0] + "\n    pass"
+
+        for name, value in members.items():
+            value_str = (
+                f"float('{repr(value)}')" if isinstance(value, float) else repr(value)
+            )
+            lines.append(f"    {name} = {value_str}")
+        return "\n".join(lines)

graph_net/paddle/naive_graph_decomposer.py

@@ -1,4 +1,7 @@
 import os
+from typing import List
+import paddle
+from graph_net import imp_util
 from graph_net.paddle.extractor import GraphExtractor as BuiltinGraphExtractor
 
 
@@ -19,47 +22,67 @@ def __init__(
 
     def make_config(
         self,
-        split_positions=(),
+        split_positions=None,
         group_head_and_tail=False,
         chain_style=False,
         output_dir="./tmp/naive_decomposer_dir",
+        post_extract_process_path=None,
+        post_extract_process_class_name=None,
+        post_extract_process_config=None,
     ):
-        for pos in split_positions:
+        assert not chain_style, "chain_style=True is not supported now."
+        if split_positions is not None:
             assert isinstance(
-                pos, int
-            ), f"split_positions should be list of int, {split_positions=}"
+                split_positions, (tuple, list)
+            ), f"split_positions is expected to be tuple or list, but recived {split_positions=}"
+            for pos in split_positions:
+                assert isinstance(
+                    pos, int
+                ), f"split_positions is expected to be tuple or list of int, but recived {split_positions=}"
         return {
             "split_positions": split_positions,
             "group_head_and_tail": group_head_and_tail,
             "chain_style": chain_style,
             "output_dir": output_dir,
+            "post_extract_process_path": post_extract_process_path,
+            "post_extract_process_class_name": post_extract_process_class_name,
+            "post_extract_process_config": post_extract_process_config,
         }
 
     def __call__(self, **input_dict):
         extracted_model = self.get_naive_decomposer_extractor()(**input_dict)
         return extracted_model
 
     def get_naive_decomposer_extractor(self):
-        return NaiveDecomposerExtractor(self)
+        return NaiveDecomposerExtractor(
+            config=self.config,
+            parent_model=self.model,
+            parent_model_name=self.name,
+            parent_input_spec=self.input_spec,
+        )
 
 
 class NaiveDecomposerExtractor:
-    def __init__(self, parent_graph_extractor):
-        super().__init__()
-        self.parent_graph_extractor = parent_graph_extractor
+    def __init__(
+        self,
+        config: dict,
+        parent_model: paddle.nn.Layer,
+        parent_model_name: str,
+        parent_input_spec: List[paddle.static.InputSpec],
+    ):
+        self.config = config
         self.extracted = False
+        self.parent_model_path = os.path.dirname(parent_model.__graph_net_file_path__)
         self.builtin_extractor = BuiltinGraphExtractor(
-            model=parent_graph_extractor.model,
-            name=parent_graph_extractor.name,
-            dynamic=parent_graph_extractor.dynamic,
-            input_spec=parent_graph_extractor.input_spec,
-            workspace_path=self.parent_graph_extractor.config["output_dir"],
+            model=parent_model,
+            name=parent_model_name,
+            dynamic=False,
+            input_spec=parent_input_spec,
+            workspace_path=self.config["output_dir"],
         )
-        self.split_positions = self.parent_graph_extractor.config["split_positions"]
-        self.group_head_and_tail = self.parent_graph_extractor.config[
-            "group_head_and_tail"
-        ]
-        self.post_process = self.make_post_process(self.parent_graph_extractor.config)
+        self.split_positions = self.config["split_positions"]
+        self.group_head_and_tail = self.config["group_head_and_tail"]
+        self.post_extract_process = self.make_post_extract_process(self.config)
 
     def do_extract(self, **input_dict):
         # 1. Run the model to dump pir programs
@@ -97,14 +120,17 @@ def __call__(self, **input_dict):
         if not self.extracted:
             extracted_model = self.do_extract(**input_dict)
             self.extracted = True
-        # if self.extracted:
-        #    for subgraph_path in self.subgraph_path_list:
-        #        self.post_process(subgraph_path)
+
+        for subgraph_path in self.subgraph_path_list:
+            self._post_extract_process(subgraph_path)
         return extracted_model
 
-    def make_post_process(self, config):
-        return None
-        # if config["post_process_path"] is None:
-        #    return None
-        # module = imp_util.load_module(config["post_process_path"])
-        # return module.PostExtractProcess(config["post_process_config"])
+    def _post_extract_process(self, subgraph_path):
+        return self.post_extract_process(subgraph_path)
+
+    def make_post_extract_process(self, config):
+        if config.get("post_extract_process_path") is None:
+            return lambda *args, **kwargs: None
+        module = imp_util.load_module(config["post_extract_process_path"])
+        cls = getattr(module, config["post_extract_process_class_name"])
+        return cls(config["post_extract_process_config"], self.parent_model_path)

graph_net/paddle/run_model.py

@@ -1,9 +1,7 @@
 import os
-import sys
 import json
 import base64
 import argparse
-from typing import Type
 
 os.environ["FLAGS_logging_pir_py_code_dir"] = "/tmp/dump"
 
@@ -16,6 +14,7 @@ def load_class_from_file(file_path: str, class_name: str):
     print(f"Load {class_name} from {file_path}")
     module = imp_util.load_module(file_path, "unnamed")
     model_class = getattr(module, class_name, None)
+    setattr(model_class, "__graph_net_file_path__", os.path.normpath(file_path))
     return model_class
 
 
@@ -26,7 +25,8 @@ def get_input_dict(model_path):
 
     state_dict = {}
     for k, v in params.items():
-        state_dict[k] = paddle.nn.parameter.Parameter(utils.replay_tensor(v), name=k)
+        name = v["original_name"] if v.get("original_name", None) else k
+        state_dict[k] = paddle.nn.parameter.Parameter(utils.replay_tensor(v), name=name)
     for k, v in inputs.items():
         state_dict[k] = utils.replay_tensor(v)
     return state_dict
@@ -83,4 +83,5 @@ def main(args):
         help="decorator configuration string",
     )
     args = parser.parse_args()
+    print(args)
     main(args=args)

graph_net/paddle/utils.py

@@ -1,11 +1,4 @@
-import re
-from collections import OrderedDict
-import uuid
-import json
-import os
-import argparse
 import importlib
-import inspect
 import ast
 import math
 import numpy as np
@@ -141,7 +134,7 @@ def convert_to_valid_number(data_type, value):
 
 def convert_meta_classes_to_tensors(file_path):
     current_device = paddle.device.get_device()
-    for name, cls in _get_classes(file_path):
+    for name, cls in get_meta_classes(file_path):
         attrs = {
             k: v
             for k, v in cls.__dict__.items()
@@ -169,10 +162,11 @@ def convert_meta_classes_to_tensors(file_path):
             },
             "data": data_value,
             "name": attrs.get("name"),
+            "original_name": attrs.get("original_name", None),
         }
 
 
-def _get_classes(file_path):
+def get_meta_classes(file_path):
     with open(file_path, "r", encoding="utf-8") as f:
         tree = ast.parse(f.read(), filename=file_path)