onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/tasks/fill_mask.py

@@ -0,0 +1,89 @@
+from typing import Any, Callable, Dict, Optional, Tuple
+import torch
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
+
+__TASK__ = "fill-mask"
+
+
+def reduce_model_config(config: Any) -> Dict[str, Any]:
+    """Reduces a model size."""
+    check_hasattr(config, "num_attention_heads", "num_hidden_layers")
+    kwargs = dict(
+        num_hidden_layers=min(config.num_hidden_layers, nhl()),
+        num_attention_heads=min(config.num_attention_heads, 4),
+    )
+    update_config(config, kwargs)
+    return kwargs
+
+
+def get_inputs(
+    model: torch.nn.Module,
+    config: Optional[Any],
+    batch_size: int,
+    sequence_length: int,
+    dummy_max_token_id: int,
+    add_second_input: int = 1,
+    **kwargs,  # unused
+):
+    """
+    Generates inputs for task ``fill-mask``.
+    Example:
+
+    ::
+
+        input_ids:T7s1x13[101,72654:A16789.23076923077],
+        token_type_ids:T7s1x13[0,0:A0.0],
+        attention_mask:T7s1x13[1,1:A1.0])
+    """
+    assert (
+        "cls_cache" not in kwargs
+    ), f"Not yet implemented for cls_cache={kwargs['cls_cache']!r}."
+    batch = "batch"
+    seq_length = "sequence_length"
+    shapes = {
+        "input_ids": {0: batch, 1: seq_length},
+        "token_type_ids": {0: batch, 1: seq_length},
+        "attention_mask": {0: batch, 1: seq_length},
+    }
+    inputs = dict(
+        input_ids=torch.randint(0, dummy_max_token_id, (batch_size, sequence_length)).to(
+            torch.int64
+        ),
+        token_type_ids=torch.zeros((batch_size, sequence_length)).to(torch.int64),
+        attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64),
+    )
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        assert (
+            add_second_input > 0
+        ), f"Not implemented for add_second_input={add_second_input}."
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            batch_size=batch_size + 1,
+            sequence_length=sequence_length + add_second_input,
+            dummy_max_token_id=dummy_max_token_id,
+            add_second_input=0,
+            **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:
+        check_hasattr(config, "vocab_size")
+    kwargs = dict(
+        batch_size=2,
+        sequence_length=30,
+        dummy_max_token_id=31999 if config is None else (config.vocab_size - 1),
+    )
+    return kwargs, get_inputs

onnx_diagnostic/tasks/image_classification.py

@@ -0,0 +1,144 @@
+from typing import Any, Callable, Dict, Optional, Tuple
+import torch
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
+
+__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=(
+            min(config.num_hidden_layers, nhl())
+            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: int = 1,
+    **kwargs,  # unused
+):
+    """
+    Generates inputs for task ``image-classification``.
+
+    :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 (
+        "cls_cache" not in kwargs
+    ), f"Not yet implemented for cls_cache={kwargs['cls_cache']!r}."
+    assert isinstance(
+        input_width, int
+    ), f"Unexpected type for input_width {type(input_width)}{config}"
+    assert isinstance(
+        input_height, 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
+        ),
+    )
+    if model.__class__.__name__ == "ViTForImageClassification":
+        inputs["interpolate_pos_encoding"] = True
+        shapes["interpolate_pos_encoding"] = None  # type: ignore[assignment]
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        assert (
+            add_second_input > 0
+        ), f"Not implemented for add_second_input={add_second_input}."
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            input_width=input_width + add_second_input,
+            input_height=input_height + add_second_input,
+            input_channels=input_channels,
+            batch_size=batch_size + 1,
+            dynamic_rope=dynamic_rope,
+            add_second_input=0,
+            **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