Gen input tensor constraints (#382)

* support checking model redundancy

* revert change of vision_model_test

* reformat python code.

* reformat bert_model_test.py and utils.py

* minor fix

* fix failed check by comparing directories after os.path.realpath()

* fix bugs in check_validate.sh

* set dynamic=False in single_device_runner.py

* reset graph hash

* add robustness code for generating input tensor constraints

* Introduce input_tensor_constraints.py using shape propagation logic.

* support dimension generalization for torch.Tensor.view and torch.Tensor.reshape

Author: lixinqi

graph_net/config/small_torch_samples_list.txt

@@ -1,3 +1,3 @@
 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/opus-mt-en-gmw
+# samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation

graph_net/constraint_util.py

@@ -6,6 +6,10 @@
 import copy
 import sys
 import os
+from contextlib import contextmanager
+import tempfile
+import shutil
+from pathlib import Path
 
 
 class UpdateInputTensorConstraints:
@@ -33,25 +37,33 @@ def _make_config(
         data_input_predicator_filepath,
         model_runnable_predicator_filepath,
         data_input_predicator_class_name="DataInputPredicator",
-        data_input_predicator_config=None,
         model_runnable_predicator_class_name="ModelRunner",
+        data_input_predicator_config=None,
         model_runnable_predicator_config=None,
+        dimension_generalizer_filepath=None,
+        dimension_generalizer_class_name="StaticToDynamic",
+        dimension_generalizer_config=None,
         model_path_prefix="",
         resume=False,
     ):
         if data_input_predicator_config is None:
             data_input_predicator_config = {}
         if model_runnable_predicator_config is None:
             model_runnable_predicator_config = {}
+        if dimension_generalizer_config is None:
+            dimension_generalizer_config = {}
         return {
+            "resume": resume,
+            "model_path_prefix": model_path_prefix,
             "data_input_predicator_filepath": data_input_predicator_filepath,
             "data_input_predicator_class_name": data_input_predicator_class_name,
             "data_input_predicator_config": data_input_predicator_config,
             "model_runnable_predicator_filepath": model_runnable_predicator_filepath,
             "model_runnable_predicator_class_name": model_runnable_predicator_class_name,
             "model_runnable_predicator_config": model_runnable_predicator_config,
-            "model_path_prefix": model_path_prefix,
-            "resume": resume,
+            "dimension_generalizer_filepath": dimension_generalizer_filepath,
+            "dimension_generalizer_class_name": dimension_generalizer_class_name,
+            "dimension_generalizer_config": dimension_generalizer_config,
         }
 
     def __call__(self, model_path):
@@ -74,17 +86,51 @@ def __call__(self, model_path):
         def data_input_predicator(input_var_name):
             return self.data_input_predicator(model_path, input_var_name)
 
-        def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
-            return self._is_dyn_dim_cstr_feasible(
-                model_path, tensor_metas, dyn_dim_cstr
-            )
+        with self._try_dimension_generalization(
+            model_path, tensor_metas
+        ) as tmp_model_path:
+
+            def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
+                return self._is_dyn_dim_cstr_feasible(
+                    tmp_model_path, tensor_metas, dyn_dim_cstr
+                )
 
-        dyn_dim_cstr = symbolize_data_input_dims(
-            dyn_dim_cstr,
-            is_data_input=data_input_predicator,
-            is_dyn_dim_cstr_feasible=is_dyn_dim_cstr_feasible,
+            dyn_dim_cstr = symbolize_data_input_dims(
+                dyn_dim_cstr,
+                is_data_input=data_input_predicator,
+                is_dyn_dim_cstr_feasible=is_dyn_dim_cstr_feasible,
+            )
+            self._save_dyn_dim_cstr(dyn_dim_cstr, model_path)
+
+    @contextmanager
+    def _try_dimension_generalization(self, 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()
+        decorator_cls = getattr(
+            load_module(self.config["dimension_generalizer_filepath"]),
+            self.config["dimension_generalizer_class_name"],
+        )
+        pass_obj = decorator_cls(self.config["dimension_generalizer_config"])(model)
+        if not pass_obj.need_rewrite():
+            yield model_path
+            return
+        from dataclasses import asdict
+
+        tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
+        graph_module = pass_obj.rewrite_with_tensor_meta_attrs_list(
+            tensor_meta_attrs_list
         )
-        self._save_dyn_dim_cstr(dyn_dim_cstr, model_path)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            shutil.copytree(Path(model_path), Path(tmp_dir), dirs_exist_ok=True)
+            pass_obj.save_graph_module(graph_module, tmp_dir)
+            shutil.copy(Path(tmp_dir) / "model.py", Path("/tmp/a.py"))
+            yield tmp_dir
+            # shutil.copytree(Path(tmp_dir), Path(model_path), dirs_exist_ok=True)
 
     def _save_dyn_dim_cstr(self, dyn_dim_cstr, model_path):
         cstr_code = dyn_dim_cstr.serialize_to_py_str()
@@ -106,7 +152,6 @@ def _is_dyn_dim_cstr_feasible(
         weight_meta_code = "\n".join(
             tensor_meta.serialize_to_py_str() for tensor_meta in tensor_metas
         )
-        import tempfile
 
         with tempfile.TemporaryDirectory() as tmpdir:
             for filename in ["graph_net.json", "model.py"]:

graph_net/imp_util.py

@@ -5,4 +5,5 @@ 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)
+    module.__graph_net_file_path__ = path
     return module

graph_net/model_path_handler.py

@@ -1,3 +1,4 @@
+import traceback
 import argparse
 import importlib.util
 from graph_net.imp_util import load_module
@@ -44,7 +45,11 @@ def main(args):
         except KeyboardInterrupt:
             sys.exit(-1)
         except Exception as e:
-            pass
+            print("--- Concise Error Message ---")
+            print(e)
+
+            print("\n--- Full Traceback ---")
+            traceback.print_exc()
 
 
 def _get_model_paths(args):

graph_net/test/batch_init_input_tensor_constraints_test.sh

@@ -11,12 +11,14 @@ 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",
         "model_runnable_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
-        "model_runnable_predicator_class_name": "ModelRunnablePredicator"
+        "model_runnable_predicator_class_name": "ModelRunnablePredicator",
+        "dimension_generalizer_filepath": "$GRAPH_NET_ROOT/torch/static_to_dynamic.py",
+        "dimension_generalizer_class_name": "StaticToDynamic"
     }
 }
 EOF

graph_net/test/dimension_generalization_test.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+
+# input model path
+config_json_str=$(cat <<EOF
+{
+    "decorator_path": "$GRAPH_NET_ROOT/torch/static_to_dynamic.py",
+    "decorator_class_name": "StaticToDynamic",
+    "decorator_config": {
+        "output_dir": ""
+    }
+}
+EOF
+)
+CONFIG=$(echo $config_json_str | base64 -w 0)
+
+python3 -m graph_net.torch.run_model --model-path $GRAPH_NET_ROOT/../samples/transformers-auto-model/opus-mt-en-gmw --decorator-config=$CONFIG

graph_net/test/shape_prop_batch_init_input_tensor_constraints_test.sh

@@ -16,11 +16,9 @@ config_json_str=$(cat <<EOF
         "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": "RunModelPredicator",
-        "model_runnable_predicator_config": {
-            "decorator_path": "$GRAPH_NET_ROOT/torch/shape_prop.py",
-            "decorator_class_name": "ShapePropagate"
-        }
+        "model_runnable_predicator_class_name": "ShapePropagatablePredicator",
+        "dimension_generalizer_filepath": "$GRAPH_NET_ROOT/torch/static_to_dynamic.py",
+        "dimension_generalizer_class_name": "StaticToDynamic"
     }
 }
 EOF

graph_net/torch/constraint_util.py

@@ -19,6 +19,23 @@ def __call__(self, model_path):
         return os.system(cmd) == 0
 
 
+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",
+        }
+        self.predicator = RunModelPredicator(decorator_config)
+
+    def __call__(self, model_path):
+        return self.predicator(model_path)
+
+
 class RunModelPredicator:
     def __init__(self, config=None):
         if config is None:

graph_net/torch/fx_graph_parse_util.py

@@ -0,0 +1,25 @@
+import torch
+
+
+def parse_sole_graph_module(module, inputs):
+    traced_module = None
+    traced_sample_inputs = None
+
+    def my_backend(gm, sample_inputs):
+        nonlocal traced_module
+        traced_module = gm
+        nonlocal traced_sample_inputs
+        traced_sample_inputs = sample_inputs
+        return gm.forward
+
+    torch.compile(module, backend=my_backend)(*inputs)
+    assert traced_module is not None
+    assert all(id(a) == id(b) for a, b in zip(inputs, traced_sample_inputs))
+    for node in traced_module.graph.nodes:
+        if node.op != "placeholder":
+            continue
+        assert node.target[:2] == "L_" or node.target[:2] == "l_", f"{node.target=}"
+        node.target = node.target[2:]
+        assert node.name[:2] == "L_" or node.name[:2] == "l_", f"{node.name=}"
+        node.name = node.name[2:]
+    return traced_module

graph_net/torch/run_model.py

@@ -17,6 +17,7 @@ def load_class_from_file(file_path: str, class_name: str) -> Type[torch.nn.Modul
     unnamed = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(unnamed)
     model_class = getattr(unnamed, class_name, None)
+    setattr(model_class, "__graph_net_file_path__", file_path)
     return model_class