add robustness code for generating input tensor constraints

Author: lixinqi

graph_net/config/small_torch_samples_list.txt

@@ -1,3 +1,2 @@
-samples/timm/resnetaa50d.d_in12k
-samples/timm/regnetx_016.pycls_in1k
-samples/timm/repghostnet_130.in1k
+samples/transformers-auto-model/opus-mt-en-gmw
+samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation

graph_net/constraint_util.py

@@ -2,6 +2,7 @@
 from graph_net.imp_util import load_module
 from graph_net.tensor_meta import TensorMeta
 from typing import Callable
+import functools
 import copy
 import sys
 import os
@@ -36,6 +37,7 @@ def _make_config(
         model_runnable_predicator_class_name="ModelRunner",
         model_runnable_predicator_config=None,
         model_path_prefix="",
+        resume=False,
     ):
         if data_input_predicator_config is None:
             data_input_predicator_config = {}
@@ -49,10 +51,23 @@ def _make_config(
             "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,
         }
 
     def __call__(self, model_path):
         model_path = os.path.join(self.config["model_path_prefix"], model_path)
+        print(f"{model_path=}")
+        cstr_path = os.path.join(model_path, "input_tensor_constraints.py")
+        if (
+            self.config["resume"]
+            and os.path.exists(cstr_path)
+            and DynamicDimConstraints.kSymbols in open(cstr_path).read()
+        ):
+            module = load_module(cstr_path)
+            symbols = getattr(module, DynamicDimConstraints.kSymbols)
+            if len(symbols) > 0:
+                return
+
         tensor_metas = self._get_tensor_metas(model_path)
         dyn_dim_cstr = make_dyn_dim_cstr_from_tensor_metas(tensor_metas)
 
@@ -111,6 +126,11 @@ def update_tensor_metas_by_dyn_dim_cstr(
     assert len(tensor_metas) == len(input_shapes)
     for i, tensor_meta in enumerate(tensor_metas):
         tensor_meta.shape = input_shapes[i]
+        if tensor_meta.data is not None:
+            assert isinstance(tensor_meta.data, (list, tuple))
+            size = functools.reduce(lambda a, b: a * b, tensor_meta.shape, 1)
+            doubled_data = [*tensor_meta.data, *tensor_meta.data]
+            tensor_meta.data = doubled_data[:size]
 
 
 def make_dyn_dim_cstr_from_tensor_metas(tensor_metas: list[TensorMeta]):
@@ -152,7 +172,11 @@ def dumpy_filter_fn(input_name, input_idx, axis, dim):
         cur_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
 
         def filter_fn(input_name, input_idx, axis, dim):
-            return is_data_input(input_name) and dim == picked_dim
+            return (
+                is_data_input(input_name)
+                and dim == picked_dim
+                and (dim > 1 or axis == 0)
+            )
 
         symbol = cur_dyn_dim_cstr.symbolize(filter_fn)
         if symbol is None:

graph_net/model_path_handler.py

@@ -39,7 +39,12 @@ def main(args):
     handler = _get_handler(args)
     for model_path in _get_model_paths(args):
         print(f"{model_path=}")
-        handler(model_path)
+        try:
+            handler(model_path)
+        except KeyboardInterrupt:
+            sys.exit(-1)
+        except Exception as e:
+            pass
 
 
 def _get_model_paths(args):

graph_net/test/batch_init_input_tensor_constraints_test.sh

@@ -11,6 +11,7 @@ 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",

graph_net/torch/constraint_util.py

@@ -7,7 +7,7 @@ def __init__(self, config):
         self.config = config
 
     def __call__(self, model_path, input_var_name: str) -> bool:
-        return not ("self_" in input_var_name and "_modules_" in input_var_name)
+        return not ("_self_" in input_var_name)
 
 
 class ModelRunnablePredicator: