Implement NaiveDecomposer for paddle.

Author: Xreki

graph_net/paddle/extractor.py

@@ -11,6 +11,34 @@
 from graph_net.paddle import utils
 
 
+# used as configuration of python -m graph_net.paddle.run_model
+class RunModelDecorator:
+    def __init__(self, config):
+        self.config = self.make_config(**config)
+
+    def __call__(self, model):
+        return extract(**self.config)(model)
+
+    def make_config(
+        self,
+        name=None,
+        dynamic=False,
+        input_spec=None,
+        custom_extractor_path: str = None,
+        custom_extractor_config: dict = None,
+    ):
+        assert name is not None
+        return {
+            "name": name,
+            "dynamic": dynamic,
+            "input_spec": input_spec,
+            "extractor_config": {
+                "custom_extractor_path": custom_extractor_path,
+                "custom_extractor_config": custom_extractor_config,
+            },
+        }
+
+
 class GraphExtractor:
     def __init__(
         self,
@@ -26,7 +54,10 @@ def __init__(
         self.input_spec = input_spec
         assert not self.dynamic, "dynamic=True is not supported now!"
 
-        self.subgraph_counter = 0
+        self.num_subgraphs = 0
+        self.num_samples_of_all_subgraphs = 0
+        self.subgraph_idx2samples = None
+
         self.dump_path = os.environ.get("GRAPH_NET_PIR_DUMP_WORKSPACE", "/tmp")
         self.workspace_path = (
             workspace_path
@@ -57,30 +88,23 @@ def prepare_to_extract(self, model_dump_path):
         )
         return old_flags
 
-    def write_to_file(self, filepath, content):
-        print(f"Write to {filepath}")
-        with open(filepath, "w") as f:
-            f.write(content)
-
-    def __call__(self, **input_dict):
-        # 1. Get model dump path
-        model_dump_path = os.path.join(self.dump_path, self.name)
-        old_flags = self.prepare_to_extract(model_dump_path)
-
+    def run_model(self, **input_dict):
         if self.input_spec is None:
             self.input_spec = [
                 paddle.static.InputSpec(value.shape, value.dtype, name=name)
                 for name, value in input_dict.items()
                 if isinstance(value, paddle.Tensor)
             ]
 
-        # 2. Run the model to dump pir programs
         static_model = paddle.jit.to_static(
             self.model, input_spec=self.input_spec, full_graph=True
         )
         static_model(**input_dict)
+        return static_model
 
-        # 3. Convert pir programs to graphnet samples
+    def translate_pir_program_to_sample_codes(
+        self, model_dump_path, split_positions=None
+    ):
         ir_programs_path = os.path.join(model_dump_path, "exec_programs.py")
         example_inputs_path = os.path.join(
             model_dump_path, "programs_example_input_tensor_meta.py"
@@ -92,29 +116,73 @@ def __call__(self, **input_dict):
             example_inputs_path
         ), f"{example_inputs_path} is not a regular file."
 
+        # Arguments for graph decomposer
+        op_example_inputs_path = (
+            os.path.join(model_dump_path, "op_example_input_tensor_meta.py")
+            if split_positions
+            else None
+        )
+        split_positions = (
+            ",".join(map(str, split_positions))
+            if split_positions and isinstance(split_positions, (tuple, list))
+            else split_positions
+        )
+
         graphnet_samples = generate_samples(
             model_name=self.name,
             ir_programs=ir_programs_path,
             example_inputs=example_inputs_path,
+            op_example_inputs=op_example_inputs_path,
+            split_positions=split_positions,
             eval_mode=True,
         )
 
+        self.subgraph_idx2samples = {}
+        for sample in graphnet_samples:
+            if sample.subgraph_idx not in self.subgraph_idx2samples.keys():
+                self.subgraph_idx2samples[sample.subgraph_idx] = []
+            self.subgraph_idx2samples[sample.subgraph_idx].append(sample)
+
+        self.num_subgraphs = len(self.subgraph_idx2samples)
+        self.num_samples_of_all_subgraphs = len(graphnet_samples)
+        return self.subgraph_idx2samples
+
+    def write_sample_to_file(self, subgraph_path, sample):
+        def write_to_file(filepath, content):
+            print(f"Write to {filepath}")
+            with open(filepath, "w") as f:
+                f.write(content)
+
+        if not os.path.exists(subgraph_path):
+            os.makedirs(subgraph_path, exist_ok=True)
+        write_to_file(f"{subgraph_path}/model.py", sample.model)
+        write_to_file(f"{subgraph_path}/weight_meta.py", sample.weight_meta)
+        write_to_file(f"{subgraph_path}/input_meta.py", sample.input_meta)
+        with open(os.path.join(subgraph_path, "graph_net.json"), "w") as f:
+            json.dump(sample.metadata, f, indent=4)
+
+    def __call__(self, **input_dict):
+        # 1. Get model dump path
+        model_dump_path = os.path.join(self.dump_path, self.name)
+        old_flags = self.prepare_to_extract(model_dump_path)
+
+        # 2. Run the model to dump pir programs
+        static_model = self.run_model(**input_dict)
+
+        # 3. Convert pir programs to graphnet samples
+        self.translate_pir_program_to_sample_codes(
+            model_dump_path, split_positions=None
+        )
+
         # 4. Save to model_path
         model_path = os.path.join(self.workspace_path, self.name)
-        self.subgraph_counter = len(graphnet_samples)
-        for i, sample in enumerate(graphnet_samples):
-            subgraph_path = (
-                model_path
-                if self.subgraph_counter == 1
-                else os.path.join(model_path, f"subgraph_{i}")
-            )
-            if not os.path.exists(subgraph_path):
-                os.makedirs(subgraph_path, exist_ok=True)
-            self.write_to_file(f"{subgraph_path}/model.py", sample.model)
-            self.write_to_file(f"{subgraph_path}/weight_meta.py", sample.weight_meta)
-            self.write_to_file(f"{subgraph_path}/input_meta.py", sample.input_meta)
-            with open(os.path.join(subgraph_path, "graph_net.json"), "w") as f:
-                json.dump(sample.metadata, f, indent=4)
+        for subgraph_idx, samples in self.subgraph_idx2samples.items():
+            assert len(samples) == 1
+            if self.num_samples_of_all_subgraphs == 1:
+                subgraph_path = model_path
+            else:
+                subgraph_path = os.path.join(model_path, f"subgraph_{subgraph_idx}")
+            self.write_sample_to_file(subgraph_path, samples[0])
 
         print(
             f"Graph and tensors for '{self.name}' extracted successfully to: {model_path}"
@@ -125,10 +193,42 @@ def __call__(self, **input_dict):
         return static_model
 
 
-def extract(name, dynamic=False, input_spec=None):
+def extract(name, dynamic=False, input_spec=None, extractor_config: dict = None):
+    """
+    Extract computation graphs from PaddlePaddle nn.Layer.
+    The extracted computation graphs will be saved into directory of env var $GRAPH_NET_EXTRACT_WORKSPACE.
+
+    Args:
+        name (str): The name of the model, used as the directory name for saving.
+        dynamic (bool): Enable dynamic shape support in paddle.jit.to_static.
+        input_spec (list[InputSpec] | tuple[InputSpec]): InputSpec for input tensors, which includes tensor's name, shape and dtype.
+            When dynamic is False, input_spec can be inferred automatically.
+
+    Returns:
+        wrapper or decorator
+    """
+
+    extractor_config = make_extractor_config(extractor_config)
+
+    def get_graph_extractor_maker():
+        custom_extractor_path = extractor_config["custom_extractor_path"]
+        custom_extractor_config = extractor_config["custom_extractor_config"]
+        if custom_extractor_path is None:
+            return GraphExtractor
+        import importlib.util as imp
+
+        print(f"Import graph_extractor from {custom_extractor_path}")
+        # import custom_extractor_path as graph_extractor
+        spec = imp.spec_from_file_location("graph_extractor", custom_extractor_path)
+        graph_extractor = imp.module_from_spec(spec)
+        spec.loader.exec_module(graph_extractor)
+        cls = graph_extractor.GraphExtractor
+        return lambda *args, **kwargs: cls(custom_extractor_config, *args, **kwargs)
+
     def wrapper(model: paddle.nn.Layer):
         assert isinstance(model, paddle.nn.Layer), f"{type(model)=}"
-        return GraphExtractor(model, name, dynamic, input_spec)
+        extractor = get_graph_extractor_maker()(model, name, dynamic, input_spec)
+        return extractor
 
     def decorator(module_class):
         def constructor(*args, **kwargs):
@@ -147,3 +247,18 @@ def decorator_or_wrapper(obj):
             )
 
     return decorator_or_wrapper
+
+
+def make_extractor_config(extractor_config):
+    kwargs = extractor_config if extractor_config is not None else {}
+    return make_extractor_config_impl(**kwargs)
+
+
+def make_extractor_config_impl(
+    custom_extractor_path: str = None, custom_extractor_config: dict = None
+):
+    config = custom_extractor_config if custom_extractor_config is not None else {}
+    return {
+        "custom_extractor_path": custom_extractor_path,
+        "custom_extractor_config": config,
+    }

graph_net/paddle/naive_graph_decomposer.py

@@ -0,0 +1,52 @@
+class GraphExtractor:
+    def __init__(
+        self,
+        config: dict,
+        model,
+        name,
+        dynamic,
+        input_spec=None,
+    ):
+        self.subgraph_counter = 0
+        self.model = model
+        self.name = name
+        self.dynamic = dynamic
+        self.input_spec = input_spec
+        self.config = self.make_config(**config)
+
+    def make_config(
+        self,
+        split_positions=(),
+        group_head_and_tail=False,
+        chain_style=False,
+        output_dir="./tmp/naive_decomposer_dir",
+        filter_path=None,
+        filter_config=None,
+    ):
+        for pos in split_positions:
+            assert isinstance(
+                pos, int
+            ), f"split_positions should be list of int, {split_positions=}"
+        return {
+            "split_positions": split_positions,
+            "group_head_and_tail": group_head_and_tail,
+            "chain_style": chain_style,
+            "output_dir": output_dir,
+            "filter_path": filter_path,
+            "filter_config": filter_config if filter_config is not None else {},
+        }
+
+    def __call__(self, **input_dict):
+        config = {
+            k: v
+            for k, v in self.config.items()
+            if k in {"split_positions", "group_head_and_tail", "chain_style"}
+        }
+        static_model = paddle.jit.to_static(
+            self.model, input_spec=self.input_spec, full_graph=True
+        )
+        static_model(**input_dict)
+        return static_model
+
+    def get_naive_decomposer_extractor(self, submodule, seq_no):
+        return NaiveDecomposerExtractor(self, submodule, seq_no)

graph_net/paddle/run_model.py

@@ -0,0 +1,83 @@
+import sys
+import json
+import base64
+import argparse
+import importlib.util
+from typing import Type
+
+import paddle
+from graph_net.paddle import utils
+
+
+def load_class_from_file(file_path: str, class_name: str):
+    print(f"Load {class_name} from {file_path}")
+    spec = importlib.util.spec_from_file_location("unnamed", file_path)
+    unnamed = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(unnamed)
+    model_class = getattr(unnamed, class_name, None)
+    return model_class
+
+
+def get_input_dict(model_path):
+    inputs_params = utils.load_converted_from_text(f"{model_path}")
+    params = inputs_params["weight_info"]
+    inputs = inputs_params["input_info"]
+
+    params.update(inputs)
+    state_dict = {k: utils.replay_tensor(v) for k, v in params.items()}
+    return state_dict
+
+
+def _convert_to_dict(config_str):
+    if config_str is None:
+        return {}
+    config_str = base64.b64decode(config_str).decode("utf-8")
+    config = json.loads(config_str)
+    assert isinstance(config, dict), f"config should be a dict. {config_str=}"
+    return config
+
+
+def _get_decorator(args):
+    if args.decorator_config is None:
+        return lambda model: model
+    decorator_config = _convert_to_dict(args.decorator_config)
+    if "decorator_path" not in decorator_config:
+        return lambda model: model
+    decorator_class = load_class_from_file(
+        decorator_config["decorator_path"], class_name="RunModelDecorator"
+    )
+    return decorator_class(decorator_config.get("decorator_config", {}))
+
+
+def main(args):
+    model_path = args.model_path
+    model_class = load_class_from_file(
+        f"{model_path}/model.py", class_name="GraphModule"
+    )
+    assert model_class is not None
+    model = model_class()
+    print(f"{model_path=}")
+
+    model = _get_decorator(args)(model)
+    input_dict = get_input_dict(args.model_path)
+
+    model(**input_dict)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="load and run model")
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        required=True,
+        help="Path to folder e.g '../../paddle_samples/PaddleX/ResNet18'",
+    )
+    parser.add_argument(
+        "--decorator-config",
+        type=str,
+        required=False,
+        default=None,
+        help="decorator configuration string",
+    )
+    args = parser.parse_args()
+    main(args=args)

graph_net/torch/extractor.py

@@ -178,7 +178,7 @@ def extract(
         dynamic (bool): Enable dynamic shape support in torch.compile.
 
     Returns:
-        wrapper or decorector
+        wrapper or decorator
 
     Examples:
         >>> # wrapper style: