Gen input tensor constraints (#380)

* 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.

Author: lixinqi

graph_net/config/small_torch_samples_list.txt

@@ -1,2 +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

graph_net/model_path_handler.py

@@ -58,6 +58,7 @@ def _get_model_paths(args):
                 for line in f
                 for clean_line in [line.strip()]
                 if len(clean_line) > 0
+                if not clean_line.startswith("#")
             )
 
 

graph_net/test/shape_prop_batch_init_input_tensor_constraints_test.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+
+# input model path
+MODEL_NAME=resnet18
+MODEL_PATH_IN_SAMPLES=/timm/$MODEL_NAME
+config_json_str=$(cat <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/constraint_util.py",
+    "handler_class_name": "UpdateInputTensorConstraints",
+    "handler_config": {
+        "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": "RunModelPredicator",
+        "model_runnable_predicator_config": {
+            "decorator_path": "$GRAPH_NET_ROOT/torch/shape_prop.py",
+            "decorator_class_name": "ShapePropagate"
+        }
+    }
+}
+EOF
+)
+CONFIG=$(echo $config_json_str | base64 -w 0)
+
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/small_torch_samples_list.txt --handler-config=$CONFIG

graph_net/torch/constraint_util.py

@@ -17,3 +17,21 @@ def __init__(self, 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
+
+
+class RunModelPredicator:
+    def __init__(self, config=None):
+        if config is None:
+            config = {}
+        self.config = config
+
+    def __call__(self, model_path):
+        import json
+        import base64
+
+        json_string = json.dumps(self.config)
+        json_bytes = json_string.encode("utf-8")
+        b64_encoded_bytes = base64.b64encode(json_bytes)
+        decorator_config = b64_encoded_bytes.decode("utf-8")
+        cmd = f"{sys.executable} -m graph_net.torch.run_model --model-path {model_path} --decorator-config {decorator_config}"
+        return os.system(cmd) == 0

graph_net/torch/run_model.py

@@ -35,8 +35,10 @@ def _get_decorator(args):
     decorator_config = _convert_to_dict(args.decorator_config)
     if "decorator_path" not in decorator_config:
         return lambda model: model
+    class_name = decorator_config.get("decorator_class_name", "RunModelDecorator")
     decorator_class = load_class_from_file(
-        decorator_config["decorator_path"], class_name="RunModelDecorator"
+        decorator_config["decorator_path"],
+        class_name=class_name,
     )
     return decorator_class(decorator_config.get("decorator_config", {}))
 

graph_net/torch/shape_prop.py

@@ -0,0 +1,30 @@
+import torch
+from typing import Union, Callable
+from torch.fx.passes.shape_prop import ShapeProp
+import inspect
+
+
+# used as configuration of python3 -m graph_net.torch.run_model
+class ShapePropagate:
+    def __init__(self, config=None):
+        if config is None:
+            config = {}
+        self.config = config
+
+    def __call__(self, module):
+        return ShapePropModule(self.config, module)
+
+
+class ShapePropModule(torch.nn.Module):
+    def __init__(self, config, module):
+        super().__init__()
+        self.config = config
+        self.module = module
+
+    def forward(self, *args, **kwargs):
+        assert len(args) == 0
+        traced_model = torch.fx.symbolic_trace(self.module)
+        inputs = [
+            kwargs[name] for name in inspect.signature(self.module.forward).parameters
+        ]
+        propagated_model = ShapeProp(traced_model).propagate(*inputs)