Enhance the performance of the input_tensor_constraints.py file generation process.

Author: lixinqi

graph_net/config/small_torch_samples_list.txt

@@ -1,3 +1,4 @@
+samples/timm/davit_giant
 # samples/timm/resnetaa50d.d_in12k
 # samples/transformers-auto-model/opus-mt-en-gmw
-samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation
+# samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation

graph_net/constraint_util.py

@@ -1,3 +1,4 @@
+import logging
 from graph_net.dynamic_dim_constraints import DynamicDimConstraints
 from contextlib import AbstractContextManager
 from graph_net.imp_util import load_module
@@ -88,14 +89,21 @@ def __call__(self, model_path):
                 return
 
         tensor_metas = self._get_tensor_metas(model_path)
+        tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
+        logging.warning(f"before create_inputs_by_metas")
+        inputs = self.get_dimension_generalizer().create_inputs_by_metas(
+            module=self.get_model(model_path),
+            tensor_meta_attrs_list=tensor_meta_attrs_list,
+        )
+        logging.warning(f"after create_inputs_by_metas")
         dyn_dim_cstr = make_dyn_dim_cstr_from_tensor_metas(tensor_metas)
 
         def data_input_predicator(input_var_name):
             return self.data_input_predicator(model_path, input_var_name)
 
         def get_tmp_model_path_ctx_mgr(dim_axes_pairs):
             return self._try_dimension_generalization(
-                dim_axes_pairs, model_path, tensor_metas
+                dim_axes_pairs, model_path, inputs
             )
 
         def get_predicator_is_dyn_dim_cstr_feasible(tmp_model_path):
@@ -128,28 +136,44 @@ def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
                     )
                     sys.exit(0)
 
-    @contextmanager
-    def _try_dimension_generalization(self, dim_axes_pairs, model_path, tensor_metas):
-        if self.config["dimension_generalizer_filepath"] is None:
-            yield model_path, ()
-            return
-        py_module = load_module(os.path.join(model_path, "model.py"))
-        GraphModule = getattr(py_module, "GraphModule")
-        GraphModule.__graph_net_file_path__ = py_module.__graph_net_file_path__
-        model = GraphModule()
+    def get_dimension_generalizer(self):
+        if hasattr(self, "_dim_generalizer"):
+            return self._dim_generalizer
+        assert self.config["dimension_generalizer_filepath"] is not None
         decorator_cls = getattr(
             load_module(self.config["dimension_generalizer_filepath"]),
             self.config["dimension_generalizer_class_name"],
         )
-        dim_generalizer = decorator_cls(self.config["dimension_generalizer_config"])
+        self._dim_generalizer = decorator_cls(
+            self.config["dimension_generalizer_config"]
+        )
+        return self._dim_generalizer
+
+    def get_model(self, model_path):
+        py_module = load_module(os.path.join(model_path, "model.py"))
+        GraphModule = getattr(py_module, "GraphModule")
+        GraphModule.__graph_net_file_path__ = py_module.__graph_net_file_path__
+        return GraphModule()
+
+    @contextmanager
+    def _try_dimension_generalization(self, dim_axes_pairs, model_path, inputs):
+        logging.warning(f"enter _try_dimension_generalization")
+        if self.config["dimension_generalizer_filepath"] is None:
+            yield model_path, ()
+            return
+        model = self.get_model(model_path)
+        dim_generalizer = self.get_dimension_generalizer()
         dim_gen_pass = dim_generalizer(model, dim_axes_pairs)
-        tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
-        inputs = dim_gen_pass.create_inputs_by_metas(tensor_meta_attrs_list)
-        if not dim_gen_pass.need_rewrite(inputs):
+        logging.warning(f"before need_rewrite")
+        need_rewrite = dim_gen_pass.need_rewrite(inputs)
+        logging.warning(f"after need_rewrite")
+        if not need_rewrite:
             yield model_path, ()
             return
 
+        logging.warning(f"before rewrite")
         graph_module = dim_gen_pass.rewrite(inputs)
+        logging.warning(f"after rewrite")
         with tempfile.TemporaryDirectory() as tmp_dir:
             shutil.copytree(Path(model_path), Path(tmp_dir), dirs_exist_ok=True)
             dim_gen_pass.save_graph_module(graph_module, tmp_dir)
@@ -300,6 +324,7 @@ def append_dim_gen_pass_names(dim_gen_pass_names):
         )
 
     for i, picked_dim in enumerate(unqiue_dims):
+        logging.warning(f"{i=} {picked_dim=}")
         cur_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
 
         def filter_fn(input_name, input_idx, axis, dim):
@@ -319,11 +344,17 @@ def filter_fn(input_name, input_idx, axis, dim):
             (dim, axes) for dim in unqiue_dims[: i + 1] for axes in [dim2axes[dim]]
         )
         ctx_mgr = dyn_dim_cstr_feasibility_ctx_mgr
+        logging.warning(f"before dyn_dim_cstr_feasibility_ctx_mgr")
         with ctx_mgr(dim_axes_pairs) as dyn_dim_cstr_feasibility:
+            logging.warning(f"enter dyn_dim_cstr_feasibility_ctx_mgr")
             tmp_dyn_dim_cstr = copy.deepcopy(cur_dyn_dim_cstr)
             tmp_dyn_dim_cstr.update_symbol2example_value(sym2example_value)
-            if not dyn_dim_cstr_feasibility(tmp_dyn_dim_cstr):
+            logging.warning(f"before dyn_dim_cstr_feasibility")
+            is_dyn_dim_cstr_feasible = dyn_dim_cstr_feasibility(tmp_dyn_dim_cstr)
+            logging.warning(f"after dyn_dim_cstr_feasibility")
+            if not is_dyn_dim_cstr_feasible:
                 continue
             dyn_dim_cstr = cur_dyn_dim_cstr
             append_dim_gen_pass_names(dyn_dim_cstr_feasibility.dim_gen_pass_names)
+            logging.warning(f"leave dyn_dim_cstr_feasibility_ctx_mgr")
     return dyn_dim_cstr, total_dim_gen_pass_names

graph_net/dynamic_dim_constraints.py

@@ -44,6 +44,8 @@ def symbolize(
         ]
         Returns created symbol.
         """
+        import logging
+
         InputDim = namedtuple("InputDim", ["input_idx", "axis", "dim"])
         input_dims = [
             InputDim(input_idx, axis, dim)

graph_net/model_path_handler.py

@@ -10,6 +10,11 @@
 import json
 import base64
 from contextlib import contextmanager
+import logging
+
+logging.basicConfig(
+    level=logging.WARNING, format="%(asctime)s [%(levelname)s] %(message)s"
+)
 
 
 def _convert_to_dict(config_str):

graph_net/test/shape_prop_batch_init_input_tensor_constraints_test.sh

@@ -11,7 +11,7 @@ config_json_str=$(cat <<EOF
     "handler_path": "$GRAPH_NET_ROOT/constraint_util.py",
     "handler_class_name": "UpdateInputTensorConstraints",
     "handler_config": {
-        "resume": true,
+        "resume": false,
         "model_path_prefix": "$GRAPH_NET_ROOT/../",
         "data_input_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
         "data_input_predicator_class_name": "NaiveDataInputPredicator",
@@ -26,4 +26,4 @@ EOF
 )
 CONFIG=$(echo $config_json_str | base64 -w 0)
 
-python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/torch_samples_list.txt --handler-config=$CONFIG
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/small_torch_samples_list.txt --handler-config=$CONFIG

graph_net/tools/batch_init_input_tensor_constraints.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+
+# input model path
+# model_runnable_predicator=ShapePropagatablePredicator
+model_runnable_predicator=ModelRunnablePredicator
+config_json_str=$(cat <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/constraint_util.py",
+    "handler_class_name": "UpdateInputTensorConstraints",
+    "handler_config": {
+        "resume": true,
+        "model_path_prefix": "$GRAPH_NET_ROOT/../",
+        "data_input_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
+        "data_input_predicator_class_name": "NaiveDataInputPredicator",
+        "model_runnable_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
+        "model_runnable_predicator_class_name": "$model_runnable_predicator",
+        "dimension_generalizer_filepath": "$GRAPH_NET_ROOT/torch/static_to_dynamic.py",
+        "dimension_generalizer_class_name": "StaticToDynamic",
+        "last_model_log_file": "/tmp/a.py"
+    }
+}
+EOF
+)
+CONFIG=$(echo $config_json_str | base64 -w 0)
+
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/torch_samples_list.txt --handler-config=$CONFIG

graph_net/torch/constraint_util.py

@@ -1,5 +1,6 @@
 import sys
 import os
+import graph_net
 
 
 class NaiveDataInputPredicator:
@@ -12,23 +13,27 @@ def __call__(self, model_path, input_var_name: str) -> bool:
 
 class ModelRunnablePredicator:
     def __init__(self, config):
-        self.config = config
+        if config is None:
+            config = {}
+
+        graph_net_root = os.path.dirname(graph_net.__file__)
+        decorator_config = {"use_dummy_inputs": True}
+        self.predicator = RunModelPredicator(decorator_config)
 
     def __call__(self, model_path):
-        cmd = f"{sys.executable} -m graph_net.torch.run_model --model-path {model_path}"
-        return os.system(cmd) == 0
+        return self.predicator(model_path)
 
 
 class ShapePropagatablePredicator:
     def __init__(self, config=None):
         if config is None:
             config = {}
-        import graph_net
 
         graph_net_root = os.path.dirname(graph_net.__file__)
         decorator_config = {
             "decorator_path": f"{graph_net_root}/torch/shape_prop.py",
             "decorator_class_name": "ShapePropagate",
+            "use_dummy_inputs": True,
         }
         self.predicator = RunModelPredicator(decorator_config)
 

graph_net/torch/fx_graph_cache_util.py

@@ -1,3 +1,4 @@
+import logging
 import torch
 import copy
 import os
@@ -13,8 +14,12 @@ def parse_immutable_model_path_into_sole_graph_module(model_path):
     if model_path not in g_model_path2graph_module:
         module = _get_torch_module(model_path)
         tensor_metas = _get_tensor_metas(model_path)
+        logging.warning("before _create_inputs_by_metas")
         inputs = _create_inputs_by_metas(module, tensor_metas)
+        logging.warning("after _create_inputs_by_metas")
+        logging.warning("before parse_sole_graph_module")
         g_model_path2graph_module[model_path] = parse_sole_graph_module(module, inputs)
+        logging.warning("after parse_sole_graph_module")
     return copy.deepcopy(g_model_path2graph_module[model_path])
 
 
@@ -34,11 +39,9 @@ def _get_tensor_metas(model_path):
 
 def _create_inputs_by_metas(module, tensor_metas):
     tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
-    from graph_net.torch.utils import convert_tensor_meta_attrs_list_to_named_tensors
+    from graph_net.torch.utils import get_dummy_named_tensors
 
-    named_tensors = convert_tensor_meta_attrs_list_to_named_tensors(
-        tensor_meta_attrs_list
-    )
+    named_tensors = get_dummy_named_tensors(tensor_meta_attrs_list)
     name2tensor = {k: v for k, v in named_tensors}
     return tuple(
         name2tensor[name] for name in inspect.signature(module.forward).parameters

graph_net/torch/run_model.py

@@ -30,10 +30,7 @@ def _convert_to_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)
+def _get_decorator(decorator_config):
     if "decorator_path" not in decorator_config:
         return lambda model: model
     class_name = decorator_config.get("decorator_class_name", "RunModelDecorator")
@@ -44,6 +41,17 @@ def _get_decorator(args):
     return decorator_class(decorator_config.get("decorator_config", {}))
 
 
+def get_flag_use_dummy_inputs(decorator_config):
+    return "use_dummy_inputs" in decorator_config
+
+
+def replay_tensor(info, use_dummy_inputs):
+    if use_dummy_inputs:
+        return utils.get_dummy_tensor(info)
+    else:
+        return utils.replay_tensor(info)
+
+
 def main(args):
     model_path = args.model_path
     model_class = load_class_from_file(
@@ -53,11 +61,15 @@ def main(args):
     model = model_class()
     print(f"{model_path=}")
 
-    model = _get_decorator(args)(model)
+    decorator_config = _convert_to_dict(args.decorator_config)
+    if "decorator_path" in args.decorator_config:
+        model = _get_decorator(decorator_config)(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()}
+    use_dummy_inputs = get_flag_use_dummy_inputs(decorator_config)
+    print(f"{use_dummy_inputs=}")
+    state_dict = {k: replay_tensor(v, use_dummy_inputs) for k, v in params.items()}
 
     model(**state_dict)
 

graph_net/torch/static_to_dynamic.py

@@ -1,6 +1,7 @@
+import logging
 import torch
 import torch.fx as fx
-from graph_net.torch.utils import convert_tensor_meta_attrs_list_to_named_tensors
+from graph_net.torch.utils import get_dummy_named_tensors
 from torch.fx.passes.shape_prop import ShapeProp
 from graph_net.torch.utils import apply_templates
 from pathlib import Path
@@ -26,6 +27,13 @@ def __init__(self, config=None):
     def __call__(self, module, dim_axes_pairs):
         return StaticToDynamicModulePass(self.config, module, dim_axes_pairs)
 
+    def create_inputs_by_metas(self, module, tensor_meta_attrs_list):
+        named_tensors = get_dummy_named_tensors(tensor_meta_attrs_list)
+        name2tensor = {k: v for k, v in named_tensors}
+        return tuple(
+            name2tensor[name] for name in inspect.signature(module.forward).parameters
+        )
+
 
 class StaticToDynamicModulePass(torch.nn.Module):
     def __init__(self, config, module, dim_axes_pairs):
@@ -56,19 +64,11 @@ def make_config(self, pass_names=()):
     def get_pass_names(self):
         return self.config["pass_names"]
 
-    def create_inputs_by_metas(self, tensor_meta_attrs_list):
-        named_tensors = convert_tensor_meta_attrs_list_to_named_tensors(
-            tensor_meta_attrs_list
-        )
-        name2tensor = {k: v for k, v in named_tensors}
-        return tuple(
-            name2tensor[name]
-            for name in inspect.signature(self.module.forward).parameters
-        )
-
     def need_rewrite(self, inputs):
         try:
+            logging.warning("before _create_fx_graph_module")
             traced_module = self._create_fx_graph_module(inputs)
+            logging.warning("after _create_fx_graph_module")
             ShapeProp(traced_module).propagate(*inputs)
         except:
             return False