onnx-diagnostic 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@ Investigates onnx models.
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.4.2"
+__version__ = "0.4.4"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -336,6 +336,10 @@ def get_parser_validate() -> ArgumentParser:
         help="drops the following inputs names, it should be a list "
         "with comma separated values",
     )
+    parser.add_argument(
+        "--subfolder",
+        help="subfolder where to find the model and the configuration",
+    )
     parser.add_argument(
         "--ortfusiontype",
         required=False,
@@ -413,6 +417,7 @@ def _cmd_validate(argv: List[Any]):
             ortfusiontype=args.ortfusiontype,
             input_options=args.iop,
             model_options=args.mop,
+            subfolder=args.subfolder,
         )
         print("")
         print("-- summary --")

onnx_diagnostic/tasks/__init__.py

@@ -6,6 +6,7 @@ from . import (
     image_classification,
     image_text_to_text,
     mixture_of_expert,
+    object_detection,
     sentence_similarity,
     text_classification,
     text_generation,
@@ -20,6 +21,7 @@ __TASKS__ = [
     image_classification,
     image_text_to_text,
     mixture_of_expert,
+    object_detection,
     sentence_similarity,
     text_classification,
     text_generation,

onnx_diagnostic/tasks/image_classification.py

@@ -7,6 +7,13 @@ __TASK__ = "image-classification"
 
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
+    if (
+        hasattr(config, "model_type")
+        and config.model_type == "timm_wrapper"
+        and not hasattr(config, "num_hidden_layers")
+    ):
+        # We cannot reduce.
+        return {}
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
@@ -82,6 +89,20 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     If the configuration is None, the function selects typical dimensions.
     """
     if config is not None:
+        if (
+            hasattr(config, "model_type")
+            and config.model_type == "timm_wrapper"
+            and not hasattr(config, "num_hidden_layers")
+        ):
+            input_size = config.pretrained_cfg["input_size"]
+            kwargs = dict(
+                batch_size=2,
+                input_width=input_size[-2],
+                input_height=input_size[-1],
+                input_channels=input_size[-3],
+            )
+            return kwargs, get_inputs
+
         check_hasattr(config, ("image_size", "architectures"), "num_channels")
     if config is not None:
         if hasattr(config, "image_size"):

onnx_diagnostic/tasks/object_detection.py

@@ -0,0 +1,123 @@
+from typing import Any, Callable, Dict, Optional, Tuple
+import torch
+from ..helpers.config_helper import update_config, check_hasattr
+
+__TASK__ = "object-detection"
+
+
+def reduce_model_config(config: Any) -> Dict[str, Any]:
+    """Reduces a model size."""
+    check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
+    kwargs = dict(
+        num_hidden_layers=(
+            min(config.num_hidden_layers, 2)
+            if hasattr(config, "num_hidden_layers")
+            else len(config.hidden_sizes)
+        )
+    )
+    update_config(config, kwargs)
+    return kwargs
+
+
+def get_inputs(
+    model: torch.nn.Module,
+    config: Optional[Any],
+    input_width: int,
+    input_height: int,
+    input_channels: int,
+    batch_size: int = 2,
+    dynamic_rope: bool = False,
+    add_second_input: bool = False,
+    **kwargs,  # unused
+):
+    """
+    Generates inputs for task ``object-detection``.
+
+    :param model: model to get the missing information
+    :param config: configuration used to generate the model
+    :param batch_size: batch size
+    :param input_channels: input channel
+    :param input_width: input width
+    :param input_height: input height
+    :return: dictionary
+    """
+    assert isinstance(
+        input_width, int
+    ), f"Unexpected type for input_width {type(input_width)}{config}"
+    assert isinstance(
+        input_width, int
+    ), f"Unexpected type for input_height {type(input_height)}{config}"
+
+    shapes = {
+        "pixel_values": {
+            0: torch.export.Dim("batch", min=1, max=1024),
+            2: "width",
+            3: "height",
+        }
+    }
+    inputs = dict(
+        pixel_values=torch.randn(batch_size, input_channels, input_width, input_height).clamp(
+            -1, 1
+        ),
+    )
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            input_width=input_width + 1,
+            input_height=input_height + 1,
+            input_channels=input_channels,
+            batch_size=batch_size + 1,
+            dynamic_rope=dynamic_rope,
+            **kwargs,
+        )["inputs"]
+    return res
+
+
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
+    """
+    Inputs kwargs.
+
+    If the configuration is None, the function selects typical dimensions.
+    """
+    if config is not None:
+        if (
+            hasattr(config, "model_type")
+            and config.model_type == "timm_wrapper"
+            and not hasattr(config, "num_hidden_layers")
+        ):
+            input_size = config.pretrained_cfg["input_size"]
+            kwargs = dict(
+                batch_size=2,
+                input_width=input_size[-2],
+                input_height=input_size[-1],
+                input_channels=input_size[-3],
+            )
+            return kwargs, get_inputs
+
+        check_hasattr(config, ("image_size", "architectures"), "num_channels")
+    if config is not None:
+        if hasattr(config, "image_size"):
+            image_size = config.image_size
+        else:
+            assert config.architectures, f"empty architecture in {config}"
+            from ..torch_models.hghub.hub_api import get_architecture_default_values
+
+            default_values = get_architecture_default_values(config.architectures[0])
+            image_size = default_values["image_size"]
+    if config is None or isinstance(image_size, int):
+        kwargs = dict(
+            batch_size=2,
+            input_width=224 if config is None else image_size,
+            input_height=224 if config is None else image_size,
+            input_channels=3 if config is None else config.num_channels,
+        )
+    else:
+        kwargs = dict(
+            batch_size=2,
+            input_width=config.image_size[0],
+            input_height=config.image_size[1],
+            input_channels=config.num_channels,
+        )
+    return kwargs, get_inputs

onnx_diagnostic/tasks/text_generation.py

@@ -19,12 +19,11 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
         ("head_dim", ("hidden_size", "num_attention_heads"), "use_mambapy"),
         "num_hidden_layers",
         ("num_key_value_heads", "num_attention_heads", "use_mambapy"),
-        "intermediate_size",
         "hidden_size",
         "vocab_size",
     )
     if config.__class__.__name__ == "FalconMambaConfig":
-        check_hasattr(config, "conv_kernel", "state_size")  # 4 and 8
+        check_hasattr(config, "conv_kernel", "state_size", "intermediate_size")  # 4 and 8
         kwargs = dict(
             num_hidden_layers=min(config.num_hidden_layers, 2),
             intermediate_size=256 if config is None else min(512, config.intermediate_size),
@@ -44,17 +43,18 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
                 if hasattr(config, "num_key_value_heads")
                 else config.num_attention_heads
             ),
-            intermediate_size=(
-                min(config.intermediate_size, 24576 // 4)
-                if config.intermediate_size % 4 == 0
-                else config.intermediate_size
-            ),
             hidden_size=(
                 min(config.hidden_size, 3072 // 4)
                 if config.hidden_size % 4 == 0
                 else config.hidden_size
             ),
         )
+        if config is None or hasattr(config, "intermediate_size"):
+            kwargs["intermediate_size"] = (
+                min(config.intermediate_size, 24576 // 4)
+                if config.intermediate_size % 4 == 0
+                else config.intermediate_size
+            )
     update_config(config, kwargs)
     return kwargs
 
@@ -228,11 +228,10 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
             "vocab_size",
             ("num_attention_heads", "use_mambapy"),
             ("num_key_value_heads", "num_attention_heads", "use_mambapy"),
-            "intermediate_size",
             "hidden_size",
         )
     if config.__class__.__name__ == "FalconMambaConfig":
-        check_hasattr(config, "conv_kernel", "state_size")  # 4 and 8
+        check_hasattr(config, "conv_kernel", "state_size", "intermediate_size")  # 4 and 8
         kwargs = dict(
             batch_size=2,
             sequence_length=30,
@@ -263,7 +262,11 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
                 if config is None
                 else _pick(config, "num_key_value_heads", "num_attention_heads")
             ),
-            intermediate_size=1024 if config is None else config.intermediate_size,
             hidden_size=512 if config is None else config.hidden_size,
         )
+        if config is None or hasattr(config, "intermediate_size"):
+            kwargs["intermediate_size"] = (
+                1024 if config is None else config.intermediate_size,
+            )
+
     return kwargs, get_inputs

onnx_diagnostic/torch_export_patches/__init__.py

@@ -1,4 +1,20 @@
 from .onnx_export_errors import (
-    bypass_export_some_errors,
+    torch_export_patches,
     register_additional_serialization_functions,
 )
+
+
+# bypass_export_some_errors is the first name given to the patches.
+bypass_export_some_errors = torch_export_patches  # type: ignore
+
+
+def register_flattening_functions(verbose: int = 0):
+    """
+    Registers functions to serialize deserialize cache or other classes
+    implemented in :epkg:`transformers` and used as inputs.
+    This is needed whenever a model must be exported through
+    :func:`torch.export.export`.
+    """
+    from .onnx_export_serialization import _register_cache_serialization
+
+    return _register_cache_serialization(verbose=verbose)

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -93,7 +93,7 @@ def register_additional_serialization_functions(
 
 
 @contextlib.contextmanager
-def bypass_export_some_errors(
+def torch_export_patches(
     patch_sympy: bool = True,
     patch_torch: bool = True,
     patch_transformers: bool = False,
@@ -145,13 +145,13 @@ def bypass_export_some_errors(
 
     ::
 
-        with bypass_export_some_errors(patch_transformers=True) as modificator:
+        with torch_export_patches(patch_transformers=True) as modificator:
             inputs = modificator(inputs)
             onx = to_onnx(..., inputs, ...)
 
     ::
 
-        with bypass_export_some_errors(patch_transformers=True) as modificator:
+        with torch_export_patches(patch_transformers=True) as modificator:
             inputs = modificator(inputs)
             onx = torch.onnx.export(..., inputs, ...)
 
@@ -159,7 +159,7 @@ def bypass_export_some_errors(
 
     ::
 
-        with bypass_export_some_errors(patch_transformers=True) as modificator:
+        with torch_export_patches(patch_transformers=True) as modificator:
             inputs = modificator(inputs)
             ep = torch.export.export(..., inputs, ...)
 
@@ -190,7 +190,7 @@ def bypass_export_some_errors(
 
         if verbose:
             print(
-                "[bypass_export_some_errors] replace torch.jit.isinstance, "
+                "[torch_export_patches] replace torch.jit.isinstance, "
                 "torch._dynamo.mark_static_address"
             )
 
@@ -210,8 +210,8 @@ def bypass_export_some_errors(
             f_sympy_name = getattr(sympy.core.numbers.IntegerConstant, "name", None)
 
             if verbose:
-                print(f"[bypass_export_some_errors] sympy.__version__={sympy.__version__!r}")
-                print("[bypass_export_some_errors] patch sympy")
+                print(f"[torch_export_patches] sympy.__version__={sympy.__version__!r}")
+                print("[torch_export_patches] patch sympy")
 
             sympy.core.numbers.IntegerConstant.name = lambda self: f"IntCst{str(self)}"
 
@@ -228,9 +228,9 @@ def bypass_export_some_errors(
             )
 
             if verbose:
-                print(f"[bypass_export_some_errors] torch.__version__={torch.__version__!r}")
-                print(f"[bypass_export_some_errors] stop_if_static={stop_if_static!r}")
-                print("[bypass_export_some_errors] patch pytorch")
+                print(f"[torch_export_patches] torch.__version__={torch.__version__!r}")
+                print(f"[torch_export_patches] stop_if_static={stop_if_static!r}")
+                print("[torch_export_patches] patch pytorch")
 
             # torch.jit.isinstance
             f_jit_isinstance = torch.jit.isinstance
@@ -252,7 +252,7 @@ def bypass_export_some_errors(
         # torch._export.non_strict_utils.produce_guards_and_solve_constraints
         if catch_constraints:
             if verbose:
-                print("[bypass_export_some_errors] modifies shape constraints")
+                print("[torch_export_patches] modifies shape constraints")
             f_produce_guards_and_solve_constraints = (
                 torch._export.non_strict_utils.produce_guards_and_solve_constraints
             )
@@ -277,22 +277,20 @@ def bypass_export_some_errors(
             ShapeEnv._log_guard_remember = ShapeEnv._log_guard
 
             if verbose:
-                print(
-                    "[bypass_export_some_errors] assert when a dynamic dimension turns static"
-                )
-                print("[bypass_export_some_errors] replaces ShapeEnv._set_replacement")
+                print("[torch_export_patches] assert when a dynamic dimension turns static")
+                print("[torch_export_patches] replaces ShapeEnv._set_replacement")
 
             f_shape_env__set_replacement = ShapeEnv._set_replacement
             ShapeEnv._set_replacement = patched_ShapeEnv._set_replacement
 
             if verbose:
-                print("[bypass_export_some_errors] replaces ShapeEnv._log_guard")
+                print("[torch_export_patches] replaces ShapeEnv._log_guard")
             f_shape_env__log_guard = ShapeEnv._log_guard
             ShapeEnv._log_guard = patched_ShapeEnv._log_guard
 
             if stop_if_static > 1:
                 if verbose:
-                    print("[bypass_export_some_errors] replaces ShapeEnv._check_frozen")
+                    print("[torch_export_patches] replaces ShapeEnv._check_frozen")
                 f_shape_env__check_frozen = ShapeEnv._check_frozen
                 ShapeEnv._check_frozen = patched_ShapeEnv._check_frozen
 
@@ -305,7 +303,7 @@ def bypass_export_some_errors(
                 import transformers
 
                 print(
-                    f"[bypass_export_some_errors] transformers.__version__="
+                    f"[torch_export_patches] transformers.__version__="
                     f"{transformers.__version__!r}"
                 )
             revert_patches_info = patch_module_or_classes(
@@ -314,7 +312,7 @@ def bypass_export_some_errors(
 
         if custom_patches:
             if verbose:
-                print("[bypass_export_some_errors] applies custom patches")
+                print("[torch_export_patches] applies custom patches")
             revert_custom_patches_info = patch_module_or_classes(
                 custom_patches, verbose=verbose
             )
@@ -326,7 +324,7 @@ def bypass_export_some_errors(
         fct_callable = replacement_before_exporting if patch_transformers else (lambda x: x)
 
         if verbose:
-            print("[bypass_export_some_errors] done patching")
+            print("[torch_export_patches] done patching")
 
         try:
             yield fct_callable
@@ -336,7 +334,7 @@ def bypass_export_some_errors(
             #######
 
             if verbose:
-                print("[bypass_export_some_errors] remove patches")
+                print("[torch_export_patches] remove patches")
 
             if patch_sympy:
                 # tracked by https://github.com/pytorch/pytorch/issues/143494
@@ -346,7 +344,7 @@ def bypass_export_some_errors(
                     delattr(sympy.core.numbers.IntegerConstant, "name")
 
                 if verbose:
-                    print("[bypass_export_some_errors] restored sympy functions")
+                    print("[torch_export_patches] restored sympy functions")
 
             #######
             # torch
@@ -362,22 +360,22 @@ def bypass_export_some_errors(
                 torch._meta_registrations._broadcast_shapes = f__broadcast_shapes
 
                 if verbose:
-                    print("[bypass_export_some_errors] restored pytorch functions")
+                    print("[torch_export_patches] restored pytorch functions")
 
             if stop_if_static:
                 if verbose:
-                    print("[bypass_export_some_errors] restored ShapeEnv._set_replacement")
+                    print("[torch_export_patches] restored ShapeEnv._set_replacement")
 
                 ShapeEnv._set_replacement = f_shape_env__set_replacement
 
                 if verbose:
-                    print("[bypass_export_some_errors] restored ShapeEnv._log_guard")
+                    print("[torch_export_patches] restored ShapeEnv._log_guard")
 
                 ShapeEnv._log_guard = f_shape_env__log_guard
 
                 if stop_if_static > 1:
                     if verbose:
-                        print("[bypass_export_some_errors] restored ShapeEnv._check_frozen")
+                        print("[torch_export_patches] restored ShapeEnv._check_frozen")
                     ShapeEnv._check_frozen = f_shape_env__check_frozen
 
             if catch_constraints:
@@ -389,11 +387,11 @@ def bypass_export_some_errors(
                     f__check_input_constraints_for_graph
                 )
                 if verbose:
-                    print("[bypass_export_some_errors] restored shape constraints")
+                    print("[torch_export_patches] restored shape constraints")
 
             if custom_patches:
                 if verbose:
-                    print("[bypass_export_some_errors] unpatch custom patches")
+                    print("[torch_export_patches] unpatch custom patches")
                 unpatch_module_or_classes(
                     custom_patches, revert_custom_patches_info, verbose=verbose
                 )
@@ -404,7 +402,7 @@ def bypass_export_some_errors(
 
             if patch_transformers:
                 if verbose:
-                    print("[bypass_export_some_errors] unpatch transformers")
+                    print("[torch_export_patches] unpatch transformers")
                 unpatch_module_or_classes(
                     patch_transformers_list, revert_patches_info, verbose=verbose
                 )