Submodule filter (#349)

Author: lixinqi

graph_net/imp_util.py

@@ -0,0 +1,8 @@
+import importlib.util as imp
+
+
+def load_module(path, name="unamed"):
+    spec = imp.spec_from_file_location(name, path)
+    module = imp.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module

graph_net/test/chain_naive_graph_decomposer_test.sh

@@ -1,20 +1,22 @@
 #!/bin/bash
-set -x
+
+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 
-# extract subgraph 0-8, 8-16
-read -r -d '' json_str <<'EOF'
+read -r -d '' extractor_config_json_str <<EOF
 {
-    "output_dir": "/tmp/naive_decompose_workspace",
-    "split_positions": [8, 16, 32],
-    "group_head_and_tail": true,
-    "chain_style": true
+    "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,
+        "chain_style": true
+    }
 }
 EOF
-CONFIG=$(echo $json_str | base64 -w 0) 
+EXTRACTOR_CONFIG=$(echo $extractor_config_json_str | base64 -w 0)
 
 mkdir -p /tmp/naive_decompose_workspace
-GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
-os.path.dirname(graph_net.__file__))")
-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 --custom-extractor-path=$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py --custom-extractor-config=$CONFIG
+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

graph_net/test/naive_graph_decomposer_test.sh

@@ -1,17 +1,23 @@
 #!/bin/bash
 
+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 
-read -r -d '' json_str <<'EOF'
+read -r -d '' extractor_config_json_str <<EOF
 {
-    "output_dir": "/tmp/naive_decompose_workspace",
-    "split_positions": [8, 32],
-    "group_head_and_tail": true
+    "custom_extractor_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
+    "custom_extractor_config": {
+        "output_dir": "/tmp/naive_decompose_workspace",
+        "split_positions": [8, 32],
+        "group_head_and_tail": true,
+        "filter_path":"$GRAPH_NET_ROOT/torch/naive_subgraph_filter.py",
+        "filter_config": {}
+    }
 }
 EOF
-CONFIG=$(echo $json_str | base64 -w 0) 
+EXTRACTOR_CONFIG=$(echo $extractor_config_json_str | base64 -w 0)
 
 mkdir -p /tmp/naive_decompose_workspace
-GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
-os.path.dirname(graph_net.__file__))")
-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 --custom-extractor-path=$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py --custom-extractor-config=$CONFIG
+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

graph_net/torch/extractor.py

@@ -139,8 +139,7 @@ def extract(
     dynamic=True,
     mut_graph_codes=None,
     placeholder_auto_rename=False,
-    custom_extractor_path: str = None,
-    custom_extractor_config: str = None,
+    extractor_config: dict = None,
 ):
     """
     Extract computation graphs from PyTorch nn.Module.
@@ -210,7 +209,11 @@ def forward(self, s0 : torch.SymInt, L_x_ : torch.Tensor):
         >>>
     """
 
+    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
@@ -247,3 +250,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/torch/naive_graph_decomposer.py

@@ -1,18 +1,17 @@
 import os
 import torch
-import json
-import base64
 import shutil
 from typing import Union, Callable
 from graph_net.torch import utils
 from graph_net.torch.decompose_util import convert_to_submodules_graph
 from graph_net.torch.extractor import GraphExtractor as BuiltinGraphExtractor
+import graph_net.imp_util as imp_util
 
 
 class GraphExtractor:
     def __init__(
         self,
-        config_str: str,
+        config: dict,
         name,
         dynamic,
         mut_graph_codes=None,
@@ -23,14 +22,16 @@ def __init__(
         self.dynamic = dynamic
         self.mut_graph_codes = mut_graph_codes
         self.placeholder_auto_rename = placeholder_auto_rename
-        self.config = self.make_config(**self.convert_to_dict(config_str))
+        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(
@@ -41,6 +42,8 @@ def make_config(
             "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, gm: torch.fx.GraphModule, sample_inputs):
@@ -59,14 +62,6 @@ def __call__(self, gm: torch.fx.GraphModule, sample_inputs):
     def get_naive_decomposer_extractor(self, submodule, seq_no):
         return NaiveDecomposerExtractor(self, submodule, seq_no)
 
-    def convert_to_dict(self, 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
-
 
 class NaiveDecomposerExtractor(torch.nn.Module):
     def __init__(self, parent_graph_extractor, submodule, seq_no):
@@ -83,9 +78,22 @@ def __init__(self, parent_graph_extractor, submodule, seq_no):
             placeholder_auto_rename=parent_graph_extractor.placeholder_auto_rename,
             workspace_path=self.parent_graph_extractor.config["output_dir"],
         )
+        self.filter = self.make_filter(self.parent_graph_extractor.config)
 
     def forward(self, *args):
         if not self.extracted:
-            self.builtin_extractor(self.submodule, args)
+            if self.need_extract(self.submodule, args):
+                self.builtin_extractor(self.submodule, args)
             self.extracted = True
         return self.submodule(*args)
+
+    def need_extract(self, gm, sample_inputs):
+        if self.filter is None:
+            return True
+        return self.filter(gm, sample_inputs)
+
+    def make_filter(self, config):
+        if config["filter_path"] is None:
+            return None
+        module = imp_util.load_module(config["filter_path"])
+        return module.GraphFilter(config["filter_config"])

graph_net/torch/naive_subgraph_filter.py

@@ -0,0 +1,7 @@
+class GraphFilter:
+    def __init__(self, config):
+        self.config = config
+
+    def __call__(self, gm, sample_inputs):
+        print(f"GraphFilter\n{gm.code}")
+        return True

graph_net/torch/single_device_runner.py

@@ -9,6 +9,8 @@
 import sys
 from graph_net.torch.extractor import extract
 import hashlib
+import json
+import base64
 from contextlib import contextmanager
 
 
@@ -20,6 +22,15 @@ def load_class_from_file(file_path: str, class_name: str) -> Type[torch.nn.Modul
     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_sha_hash(content):
     m = hashlib.sha256()
     m.update(content.encode())
@@ -63,8 +74,7 @@ def main(args):
             kwargs = dict(
                 name=args.extract_name,
                 dynamic=False,
-                custom_extractor_path=args.custom_extractor_path,
-                custom_extractor_config=args.custom_extractor_config,
+                extractor_config=convert_to_dict(args.extractor_config),
                 **dump_graph_options,
             )
             model = extract(**kwargs)(model)
@@ -118,18 +128,11 @@ def main(args):
         help="Extracted graph's name",
     )
     parser.add_argument(
-        "--custom-extractor-path",
-        type=str,
-        required=False,
-        default=None,
-        help="Custom extractor python file path",
-    )
-    parser.add_argument(
-        "--custom-extractor-config",
+        "--extractor-config",
         type=str,
         required=False,
         default=None,
-        help="Custom extractor configuration string",
+        help="extractor configuration string",
     )
     args = parser.parse_args()
     main(args=args)