add a script to check fusability of a given model

Author: roll-away

graph_net/test/dimension_generalization_test.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+
+# input model path
+MODEL_NAME=resnet18d.ra2_in1k
+MODEL_PATH_IN_SAMPLES=/timm/$MODEL_NAME
+checker_config_json_str=$(cat <<EOF
+{
+    "post_extract_process_config": {
+        "post_extract_process_path":"$GRAPH_NET_ROOT/torch/post_extract_process_count_kernels.py",
+        "post_extract_process_class_name": "GraphFullyFusionable"
+    }
+}
+EOF
+)
+CHECKER_CONFIG=$(echo $checker_config_json_str | base64 -w 0)
+
+python3 -m graph_net.torch.check_model_fusability --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --checker-config=$CHECKER_CONFIG

graph_net/test/naive_check_if_fully_fusable.sh

@@ -20,7 +20,7 @@ decorator_config_json_str=$(cat <<EOF
             "filter_path":"$GRAPH_NET_ROOT/torch/naive_subgraph_filter.py",
             "filter_config": {},
             "post_extract_process_path":"$GRAPH_NET_ROOT/torch/post_extract_process_count_kernels.py",
-            "post_extract_process_class_name": "PostExtractProcess"
+            "post_extract_process_class_name": "GraphFullyFusionable"
         }
     }
 }

graph_net/test/naive_decomposer_and_post_extract_process_test.sh

@@ -0,0 +1,71 @@
+from . import utils
+import argparse
+import importlib.util
+import inspect
+from graph_net.imp_util import load_module
+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, class_name):
+    module = load_module(file_path)
+    return getattr(module, class_name)
+
+
+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_checker(args):
+    if args.checker_config is None:
+        return lambda model_path: model_path
+    checker_config = _convert_to_dict(args.checker_config).get(
+        "post_extract_process_config"
+    )
+    checker_class = _load_class_from_file(
+        checker_config["post_extract_process_path"],
+        class_name=checker_config["post_extract_process_class_name"],
+    )
+    return checker_class(checker_config.get("checker_config", {}))
+
+
+def main(args):
+    checker = _get_checker(args)
+    model_path = args.model_path
+    print(f"{model_path=}")
+    try:
+        checker(model_path)
+    except KeyboardInterrupt:
+        sys.exit(-1)
+    except Exception as e:
+        print(e)
+
+
+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(
+        "--checker-config",
+        type=str,
+        required=False,
+        default=None,
+        help="checker configuration string",
+    )
+    args = parser.parse_args()
+    main(args=args)

graph_net/torch/check_model_fusability.py

@@ -117,4 +117,4 @@ def make_post_extract_process(self, config):
         if config["post_extract_process_path"] is None:
             return None
         module = imp_util.load_module(config["post_extract_process_path"])
-        return module.PostExtractProcess(config["post_extract_process_path"])
+        return module.GraphFullyFusionable(config["post_extract_process_path"])

graph_net/torch/naive_graph_decomposer.py

@@ -6,7 +6,7 @@
 from torch.profiler import profile, record_function, ProfilerActivity
 
 
-class PostExtractProcess:
+class GraphFullyFusionable:
     def __init__(self, config):
         self.config = config