Merge branch 'PaddlePaddle:develop' into develop

Author: roll-away

docs/torch_to_paddle_conversion_design.md

@@ -1,3 +1,3 @@
-samples/timm/resnetaa50d.d_in12k
+# 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/config/small_torch_samples_list.txt

@@ -1,11 +1,17 @@
 from graph_net.dynamic_dim_constraints import DynamicDimConstraints
+from contextlib import AbstractContextManager
 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
+from contextlib import contextmanager
+import tempfile
+import shutil
+from pathlib import Path
+import json
 
 
 class UpdateInputTensorConstraints:
@@ -17,6 +23,7 @@ def __init__(self, config=None):
         self.model_runnable_predicator = self._make_model_runnable_predicator(
             self.config
         )
+        self.num_successful_handled_models = 0
 
     def _make_data_input_predicator(self, config):
         module = load_module(config["data_input_predicator_filepath"])
@@ -33,25 +40,37 @@ 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,
+        last_model_log_file=None,
+        limits_successfully_handled_models=None,
     ):
         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,
+            "last_model_log_file": last_model_log_file,
+            "limits_successfully_handled_models": limits_successfully_handled_models,
         }
 
     def __call__(self, model_path):
@@ -74,17 +93,80 @@ 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
+        def get_tmp_model_path_ctx_mgr(dim_axes_pairs):
+            return self._try_dimension_generalization(
+                dim_axes_pairs, model_path, tensor_metas
             )
 
-        dyn_dim_cstr = symbolize_data_input_dims(
+        def get_predicator_is_dyn_dim_cstr_feasible(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
+                )
+
+            return is_dyn_dim_cstr_feasible
+
+        dyn_dim_cstr_feasibility_ctx_mgr = DynDimCstrFeasibilityContextManager(
+            get_tmp_model_path_ctx_mgr=get_tmp_model_path_ctx_mgr,
+            get_predicator_is_dyn_dim_cstr_feasible=get_predicator_is_dyn_dim_cstr_feasible,
+        )
+        dyn_dim_cstr, dim_gen_pass_names = 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_feasibility_ctx_mgr=dyn_dim_cstr_feasibility_ctx_mgr,
         )
         self._save_dyn_dim_cstr(dyn_dim_cstr, model_path)
+        self._save_dim_gen_pass_names(dim_gen_pass_names, model_path)
+        if len(dyn_dim_cstr.symbols) > 0:
+            self.num_successful_handled_models += 1
+            limits = self.config["limits_successfully_handled_models"]
+            if limits is not None:
+                if self.num_successful_handled_models > limits:
+                    print(
+                        "`num_successful_handled_models` exceeds config `limits_successfully_handled_models`",
+                        file=sys.stderr,
+                    )
+                    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()
+        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"])
+        dim_gen_pass = dim_generalizer(model, dim_axes_pairs)
+        if not dim_gen_pass.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 = dim_gen_pass.rewrite_with_tensor_meta_attrs_list(
+            tensor_meta_attrs_list=tensor_meta_attrs_list,
+        )
+        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)
+            if self.config["last_model_log_file"] is not None:
+                log_file = Path(self.config["last_model_log_file"])
+                shutil.copy(Path(tmp_dir) / "model.py", log_file)
+            yield tmp_dir, dim_gen_pass.get_pass_names()
+
+    def _save_dim_gen_pass_names(self, dim_gen_pass_names, model_path):
+        from graph_net.graph_net_json_file_util import kDimensionGeneralizationPasses
+
+        graph_net_json_file_path = Path(f"{model_path}/graph_net.json")
+        graph_net_json = json.loads(graph_net_json_file_path.read_text())
+        graph_net_json[kDimensionGeneralizationPasses] = list(dim_gen_pass_names)
+        graph_net_json_file_path.write_text(json.dumps(graph_net_json))
 
     def _save_dyn_dim_cstr(self, dyn_dim_cstr, model_path):
         cstr_code = dyn_dim_cstr.serialize_to_py_str()
@@ -106,7 +188,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"]:
@@ -145,30 +226,82 @@ def make_dyn_dim_cstr_from_tensor_metas(tensor_metas: list[TensorMeta]):
     )
 
 
+class DynDimCstrFeasibilityPredicator:
+    def __init__(
+        self,
+        is_dyn_dim_cstr_feasible: Callable[[DynamicDimConstraints], bool],
+        dim_gen_pass_names: tuple[str],
+    ):
+        self.is_dyn_dim_cstr_feasible = is_dyn_dim_cstr_feasible
+        self.dim_gen_pass_names = dim_gen_pass_names
+
+    def __call__(self, dyn_dim_cstr: DynamicDimConstraints) -> bool:
+        return self.is_dyn_dim_cstr_feasible(dyn_dim_cstr)
+
+
+class DynDimCstrFeasibilityContextManager:
+    def __init__(
+        self,
+        get_tmp_model_path_ctx_mgr,
+        get_predicator_is_dyn_dim_cstr_feasible,
+    ):
+        self.get_tmp_model_path_ctx_mgr = get_tmp_model_path_ctx_mgr
+        self.get_predicator_is_dyn_dim_cstr_feasible = (
+            get_predicator_is_dyn_dim_cstr_feasible
+        )
+
+    @contextmanager
+    def __call__(
+        self, dim_axes_pairs
+    ) -> AbstractContextManager[DynDimCstrFeasibilityPredicator]:
+        ctx_mgr = self.get_tmp_model_path_ctx_mgr
+        with ctx_mgr(dim_axes_pairs) as (tmp_model_apth, dg_pass_names):
+            predicator = self.get_predicator_is_dyn_dim_cstr_feasible(tmp_model_apth)
+            yield DynDimCstrFeasibilityPredicator(predicator, dg_pass_names)
+
+
 def symbolize_data_input_dims(
     dyn_dim_cstr: DynamicDimConstraints,
     is_data_input: Callable[[str], bool],
-    is_dyn_dim_cstr_feasible: Callable[[DynamicDimConstraints], bool],
-) -> DynamicDimConstraints | None:
+    dyn_dim_cstr_feasibility_ctx_mgr: DynDimCstrFeasibilityContextManager,
+) -> (DynamicDimConstraints | None, tuple[str]):
     """
     is_data_input: Callable[["input_var_name:str"], bool]
     Symbolizes data input dimensions as much as possible.
     Returns new DynamicDimConstraints if success.
     Returns None if no symbolicable dim .
     """
     unqiue_dims = []
+    dim2axes = {}
 
     def dumpy_filter_fn(input_name, input_idx, axis, dim):
         if is_data_input(input_name):
             print("data_input", input_name, input_idx, axis, dim)
             if dim not in unqiue_dims:
                 unqiue_dims.append(dim)
-        # No symbolization because of returning True
+                dim2axes[dim] = []
+            dim2axes[dim].append(axis)
+        # No symbolization by returning False
         return False
 
     # Collect input dimensions into `unqiue_dims`
     assert dyn_dim_cstr.symbolize(dumpy_filter_fn) is None
-    for picked_dim in unqiue_dims:
+    total_dim_gen_pass_names = ()
+
+    def append_dim_gen_pass_names(dim_gen_pass_names):
+        nonlocal total_dim_gen_pass_names
+        total_dim_gen_pass_names = tuple(
+            [
+                *total_dim_gen_pass_names,
+                *(
+                    pass_name
+                    for pass_name in dim_gen_pass_names
+                    if pass_name not in total_dim_gen_pass_names
+                ),
+            ]
+        )
+
+    for i, picked_dim in enumerate(unqiue_dims):
         cur_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
 
         def filter_fn(input_name, input_idx, axis, dim):
@@ -184,9 +317,15 @@ def filter_fn(input_name, input_idx, axis, dim):
         sym2example_value = {symbol: picked_dim + 1}
         if not cur_dyn_dim_cstr.check_delta_symbol2example_value(sym2example_value):
             continue
-        tmp_dyn_dim_cstr = copy.deepcopy(cur_dyn_dim_cstr)
-        tmp_dyn_dim_cstr.update_symbol2example_value(sym2example_value)
-        if not is_dyn_dim_cstr_feasible(tmp_dyn_dim_cstr):
-            continue
-        dyn_dim_cstr = cur_dyn_dim_cstr
-    return dyn_dim_cstr
+        dim_axes_pairs = tuple(
+            (dim, axes) for dim in unqiue_dims[: i + 1] for axes in [dim2axes[dim]]
+        )
+        ctx_mgr = dyn_dim_cstr_feasibility_ctx_mgr
+        with ctx_mgr(dim_axes_pairs) as dyn_dim_cstr_feasibility:
+            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):
+                continue
+            dyn_dim_cstr = cur_dyn_dim_cstr
+            append_dim_gen_pass_names(dyn_dim_cstr_feasibility.dim_gen_pass_names)
+    return dyn_dim_cstr, total_dim_gen_pass_names

graph_net/constraint_util.py

@@ -0,0 +1 @@
+kDimensionGeneralizationPasses = "dimension_generalization_passes"

graph_net/graph_net_json_file_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/imp_util.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):