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

Author: Xreki

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,66 @@ 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.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 +119,19 @@ 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):
+    def _post_extract_process(self, model_path):
         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"])
+        # model_path = os.path.join(self.config["output_dir"], self.model_name)
+        # return self.post_extract_process(model_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"])

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"
 
@@ -26,7 +24,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 +82,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
@@ -169,6 +162,7 @@ def convert_meta_classes_to_tensors(file_path):
             },
             "data": data_value,
             "name": attrs.get("name"),
+            "original_name": attrs.get("original_name", None),
         }