add run_model.py to improve reusability.

Author: lixinqi

graph_net/test/naive_graph_decomposer_test.sh

@@ -4,21 +4,26 @@ GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
 os.path.dirname(graph_net.__file__))")
 
 # input model path
-MODEL_PATH_IN_SAMPLES=/timm/resnet18 
-extractor_config_json_str=$(cat <<EOF
+MODEL_NAME=resnet18
+MODEL_PATH_IN_SAMPLES=/timm/$MODEL_NAME
+decorator_config_json_str=$(cat <<EOF
 {
-    "custom_extractor_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
-    "custom_extractor_config": {
-        "output_dir": "/tmp/naive_decompose_workspace",
-        "split_positions": [8, 16, 32],
-        "group_head_and_tail": true,
-        "filter_path":"$GRAPH_NET_ROOT/torch/naive_subgraph_filter.py",
-        "filter_config": {}
+    "decorator_path": "$GRAPH_NET_ROOT/torch/extractor.py",
+    "decorator_config": {
+        "name": "$MODEL_NAME",
+        "custom_extractor_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
+        "custom_extractor_config": {
+            "output_dir": "/tmp/naive_decompose_workspace",
+            "split_positions": [8, 16, 32],
+            "group_head_and_tail": true,
+            "filter_path":"$GRAPH_NET_ROOT/torch/naive_subgraph_filter.py",
+            "filter_config": {}
+        }
     }
 }
 EOF
 )
-EXTRACTOR_CONFIG=$(echo $extractor_config_json_str | base64 -w 0)
+DECORATOR_CONFIG=$(echo $decorator_config_json_str | base64 -w 0)
 
 mkdir -p /tmp/naive_decompose_workspace
-python3 -m graph_net.torch.single_device_runner --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --enable-extract True --extract-name resnet18 --dump-graph-hash-key --extractor-config=$EXTRACTOR_CONFIG
+python3 -m graph_net.torch.run_model --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --decorator-config=$DECORATOR_CONFIG

graph_net/torch/extractor.py

@@ -3,8 +3,8 @@
 import json
 import shutil
 from typing import Union, Callable
-from . import utils
-from .fx_graph_serialize_util import serialize_graph_module_to_str
+from graph_net.torch import utils
+from graph_net.torch.fx_graph_serialize_util import serialize_graph_module_to_str
 
 torch._dynamo.config.capture_scalar_outputs = True
 torch._dynamo.config.capture_dynamic_output_shape_ops = True
@@ -13,6 +13,34 @@
 torch._dynamo.config.allow_rnn = True
 
 
+# used as configuration of python3 -m graph_net.torch.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,
+        placeholder_auto_rename=False,
+        custom_extractor_path: str = None,
+        custom_extractor_config: dict = None,
+    ):
+        assert name is not None
+        return {
+            "name": name,
+            "dynamic": dynamic,
+            "placeholder_auto_rename": placeholder_auto_rename,
+            "extractor_config": {
+                "custom_extractor_path": custom_extractor_path,
+                "custom_extractor_config": custom_extractor_config,
+            },
+        }
+
+
 class GraphExtractor:
     def __init__(
         self,

graph_net/torch/run_model.py

@@ -0,0 +1,78 @@
+from . import utils
+import argparse
+import importlib.util
+import inspect
+import torch
+import logging
+from pathlib import Path
+from typing import Type, Any
+import sys
+import json
+import base64
+from contextlib import contextmanager
+
+
+def load_class_from_file(file_path: str, class_name: str) -> Type[torch.nn.Module]:
+    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 _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)
+
+    inputs_params = utils.load_converted_from_text(f"{model_path}")
+    params = inputs_params["weight_info"]
+    state_dict = {k: utils.replay_tensor(v) for k, v in params.items()}
+
+    model(**state_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 '../../samples/torch/resnet18'",
+    )
+    parser.add_argument(
+        "--decorator-config",
+        type=str,
+        required=False,
+        default=None,
+        help="decorator configuration string",
+    )
+    args = parser.parse_args()
+    main(args=args)