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=True,
+        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,
@@ -125,10 +153,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 +207,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)