onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/reference/torch_ops/shape_ops.py

@@ -0,0 +1,121 @@
+from typing import Optional, Tuple
+import onnx
+import torch
+from . import OpRunKernel, OpRunTensor
+
+
+class ConstantOfShape_9(OpRunKernel):
+    "ConstantOfShape"
+
+    @classmethod
+    def device_dependent(cls) -> bool:
+        """
+        Returns True if the kernel needs a device to be efficiently initialized.
+        """
+        return True
+
+    def __init__(
+        self,
+        node: onnx.NodeProto,
+        version: Optional[int] = None,
+        device: Optional[torch.device] = None,
+        verbose: int = 0,
+    ):
+        super().__init__(node, version, verbose=verbose)
+        value = self.get_attribute_tensor(node, "value")
+        if value is None:
+            value = torch.tensor([0], dtype=torch.float32)
+        self.dtype = value.dtype
+        self.device = device
+        self.value = value[0]
+
+    def run(self, shape: OpRunTensor) -> OpRunTensor:
+        # The device is unknown as shapes usually take place on CPU.
+        return OpRunTensor(
+            torch.full(
+                shape.as_tuple_int, fill_value=self.value, dtype=self.dtype, device=self.device
+            )
+        )
+
+
+class Expand_8(OpRunKernel):
+    "Expand"
+
+    def run(self, data: OpRunTensor, shape: OpRunTensor) -> OpRunTensor:
+        ishape = tuple(-1 if i == 1 else i for i in shape.as_tuple_int)
+        return OpRunTensor(data.tensor.expand(ishape))
+
+
+class Reshape_14(OpRunKernel):
+    "Reshape"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.allowzero = self.get_attribute_int(node, "allowzero", 0)
+
+    def run(self, data: OpRunTensor, shape: OpRunTensor) -> OpRunTensor:
+        ishape = shape.as_tuple_int
+        assert ishape is not None, f"Unexpected return for shape={shape!r}"
+        if not self.allowzero and 0 in ishape:
+            xshape = data.tensor.shape
+            new_shape = []
+            for i, s in enumerate(ishape):
+                new_shape.append(xshape[i] if s == 0 else s)
+            return OpRunTensor(data.tensor.reshape(new_shape))
+        return OpRunTensor(data.tensor.reshape(ishape))
+
+
+class Shape_15(OpRunKernel):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.start = self.get_attribute_int(node, "start", 0)
+        self.end = self.get_attribute_int(node, "end", None)
+
+    def run(self, data: OpRunTensor) -> OpRunTensor:
+        shape = data.shape
+        sh = shape[self.start :] if self.end is None else shape[self.start : self.end]
+        return OpRunTensor(torch.tensor(sh, dtype=torch.int64), is_constant=True)
+
+
+class Split_18(OpRunKernel):
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None, verbose: int = 0):
+        super().__init__(node, version, verbose=verbose)
+        self.axis = self.get_attribute_int(node, "axis", 0)
+        self.num_outputs = self.get_attribute_int(node, "num_outputs", None)
+
+    def run(
+        self, data: OpRunTensor, split: Optional[OpRunTensor] = None
+    ) -> Tuple[OpRunTensor, ...]:
+        if split is None:
+            assert isinstance(
+                self.num_outputs, int
+            ), f"Incompatibilities: split is None and num_outputs={self.num_outputs}"
+            size = data.tensor.shape[self.axis]
+            split_size = (
+                size // self.num_outputs
+                if size % self.num_outputs == 0
+                else size // self.num_outputs + 1
+            )
+            spl = torch.split(data.tensor, split_size, dim=self.axis)
+        else:
+            spl = torch.split(data.tensor, split.as_tuple_int, dim=self.axis)
+        return tuple(OpRunTensor(t) for t in spl)
+
+
+class Squeeze_13(OpRunKernel):
+    "Squeeze"
+
+    def run(self, data: OpRunTensor, axes: Optional[OpRunTensor] = None) -> OpRunTensor:
+        if axes is None:
+            return OpRunTensor(data.tensor.squeeze())
+        return OpRunTensor(data.tensor.squeeze(axes.as_tuple_int))
+
+
+class Unsqueeze_13(OpRunKernel):
+    "Unsqueeze"
+
+    def run(self, data: OpRunTensor, axes: OpRunTensor) -> OpRunTensor:
+        t = data.tensor
+        for i in axes.as_tuple_int:
+            t = t.unsqueeze(i)
+        return OpRunTensor(t)

onnx_diagnostic/reference/torch_ops/unary_ops.py

@@ -0,0 +1,93 @@
+import torch
+from . import OpRunKernel, OpRunTensor
+
+
+class Abs_1(OpRunKernel):
+    """Abs"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.abs(x.tensor))
+
+
+class Cos_1(OpRunKernel):
+    """Cos"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.cos())
+
+
+class Erf_9(OpRunKernel):
+    """Erf"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.erf())
+
+
+class Exp_1(OpRunKernel):
+    """Exp"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.exp())
+
+
+class Identity_1(OpRunKernel):
+    "Identity"
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor)
+
+
+class IsNaN_9(OpRunKernel):
+    """IsNaN"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.isnan())
+
+
+class Log_1(OpRunKernel):
+    """Log"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.log())
+
+
+class Neg_1(OpRunKernel):
+    """Neg"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(-x.tensor)
+
+
+class Not_1(OpRunKernel):
+    """Not"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(~x.tensor)
+
+
+class Reciprocal_1(OpRunKernel):
+    """REciprocal"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(1 / x.tensor)
+
+
+class Sigmoid_6(OpRunKernel):
+    """Sqrt"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.sigmoid(x.tensor))
+
+
+class Sin_1(OpRunKernel):
+    """Sin"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.sin())
+
+
+class Sqrt_1(OpRunKernel):
+    """Sqrt"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.sqrt())

onnx_diagnostic/tasks/__init__.py

@@ -0,0 +1,90 @@
+from typing import Any, Callable, Dict, List, Tuple
+from . import (
+    automatic_speech_recognition,
+    feature_extraction,
+    fill_mask,
+    image_classification,
+    image_text_to_text,
+    image_to_video,
+    mask_generation,
+    mixture_of_expert,
+    object_detection,
+    sentence_similarity,
+    summarization,
+    text_classification,
+    text_generation,
+    text_to_image,
+    text2text_generation,
+    zero_shot_image_classification,
+)
+
+__TASKS__ = [
+    automatic_speech_recognition,
+    feature_extraction,
+    fill_mask,
+    image_classification,
+    image_text_to_text,
+    image_to_video,
+    mask_generation,
+    mixture_of_expert,
+    object_detection,
+    sentence_similarity,
+    summarization,
+    text_classification,
+    text_generation,
+    text_to_image,
+    text2text_generation,
+    zero_shot_image_classification,
+]
+
+
+def supported_tasks() -> List[str]:
+    "Returns the list of supported tasks."
+    return sorted(mod.__TASK__ for mod in __TASKS__)
+
+
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    head_size0 = (
+        config.head_dim
+        if hasattr(config, "head_dim") and config.head_dim
+        else (
+            config.hidden_size // config.num_attention_heads
+            if hasattr(config, "hidden_size") and hasattr(config, "num_attention_heads")
+            else None
+        )
+    )
+    tasks = {mod.__TASK__: mod.reduce_model_config for mod in __TASKS__}
+    assert task in tasks, f"Task {task!r} not found in {sorted(tasks)}"
+    res = tasks[task](config)
+    if head_size0 and "head_dim" in res:
+        head_size = (
+            config.head_dim
+            if hasattr(config, "head_dim") and config.head_dim
+            else config.hidden_size // config.num_attention_heads
+        )
+        assert head_size0 == head_size or head_size % 16 == 0, (
+            f"head_size should be a multiple of 16 "
+            f"(head_size0={head_size0}), res={res}, "
+            f"config=\n{config}"
+        )
+    return res
+
+
+def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+    """
+    Inputs kwargs.
+    If the configuration is None, the function selects typical dimensions.
+    It returns parameters and a function. The function creates dummy inputs
+    if it receives the parameters returned as a first result.
+
+    .. code-block:: python
+
+        config = get_pretrained_config(model_id)
+        task = task = task_from_id(name)
+        kwargs, fct = random_input_kwargs(config, task)
+        res = fct(model, config, add_second_input=False, **kwargs)
+    """
+    tasks = {mod.__TASK__: mod.random_input_kwargs for mod in __TASKS__}
+    assert task in tasks, f"Task {task!r} not found in {sorted(tasks)}"
+    return tasks[task](config)

onnx_diagnostic/tasks/automatic_speech_recognition.py

@@ -0,0 +1,188 @@
+from typing import Any, Callable, Dict, Optional, Tuple
+import torch
+import transformers
+from ..helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
+from ..helpers.config_helper import (
+    update_config,
+    check_hasattr,
+    default_num_hidden_layers as nhl,
+)
+
+__TASK__ = "automatic-speech-recognition"
+
+
+def reduce_model_config(config: Any) -> Dict[str, Any]:
+    """Reduces a model size."""
+    kwargs: Dict[str, Any] = {}
+    if hasattr(config, "num_decoder_layers"):
+        config.num_decoder_layers = min(config.num_decoder_layers, 2)
+    if hasattr(config, "decoder_layers"):
+        config.decoder_layers = min(config.decoder_layers, 2)
+    if hasattr(config, "num_hidden_layers"):
+        config.num_hidden_layers = min(config.num_hidden_layers, nhl())
+    update_config(config, kwargs)
+    return kwargs
+
+
+def get_inputs(
+    model: torch.nn.Module,
+    config: Optional[Any],
+    dummy_max_token_id: int,
+    max_source_positions: int,
+    d_model: int,
+    num_hidden_layers: int,
+    encoder_attention_heads: int,
+    encoder_layers: int,
+    decoder_layers: int,
+    head_dim: int,
+    batch_size: int = 2,
+    sequence_length: int = 30,
+    add_second_input: int = 1,
+    **kwargs,  # unused
+):
+    """
+    Generates inputs for task ``automatic-speech-recognition``.
+    Example:
+
+    ::
+
+        dict(
+            cache_position:T7s4,
+            past_key_values:EncoderDecoderCache(
+                self_attention_cache=DynamicCache[serialized](#2[#0[],#0[]]),
+                cross_attention_cache=DynamicCache[serialized](#2[#0[],#0[]])
+            ),
+            decoder_input_ids:T7s1x4,
+            encoder_outputs:BaseModelOutput(last_hidden_state:T1s1x1500x384),
+            use_cache:bool,return_dict:bool
+        )
+        dict(
+            cache_position:T7s1,
+            past_key_values:EncoderDecoderCache(
+                self_attention_cache=DynamicCache[serialized](#2[
+                    #4[T1s1x6x4x64,T1s1x6x4x64,T1s1x6x4x64,T1s1x6x4x64],
+                    #4[T1s1x6x4x64,T1s1x6x4x64,T1s1x6x4x64,T1s1x6x4x64]
+                ]),
+                cross_attention_cache=DynamicCache[serialized](#2[
+                    #4[T1s1x6x1500x64,T1s1x6x1500x64,T1s1x6x1500x64,T1s1x6x1500x64],
+                    #4[T1s1x6x1500x64,T1s1x6x1500x64,T1s1x6x1500x64,T1s1x6x1500x64]
+                ]),
+            ),
+            decoder_input_ids:T7s1x1,
+            encoder_outputs:BaseModelOutput(last_hidden_state:T1s1x1500x384),
+            use_cache:bool,return_dict:bool
+        )
+    """
+    assert (
+        "cls_cache" not in kwargs
+    ), f"Not yet implemented for cls_cache={kwargs['cls_cache']!r}."
+    batch = "batch"
+    seq_length = "seq_length"
+
+    shapes = {
+        "decoder_input_ids": {0: batch, 1: seq_length},
+        "cache_position": {0: seq_length},
+        "encoder_outputs": [{0: batch}],  # last_hidden_state
+        "past_key_values": [
+            [{0: batch} for _ in range(num_hidden_layers * 2)],
+            [{0: batch} for _ in range(num_hidden_layers * 2)],
+        ],
+    }
+    inputs = dict(
+        decoder_input_ids=torch.randint(
+            0, dummy_max_token_id, (batch_size, sequence_length)
+        ).to(torch.int64),
+        cache_position=(torch.arange(sequence_length) + 5).to(torch.int64),
+        encoder_outputs=transformers.modeling_outputs.BaseModelOutput(
+            last_hidden_state=torch.randn(batch_size, max_source_positions, d_model)
+        ),
+        past_key_values=make_encoder_decoder_cache(
+            make_dynamic_cache(
+                [
+                    (
+                        torch.randn(
+                            batch_size, encoder_attention_heads, encoder_layers, head_dim
+                        ),
+                        torch.randn(
+                            batch_size, encoder_attention_heads, encoder_layers, head_dim
+                        ),
+                    )
+                    for i in range(num_hidden_layers)
+                ]
+            ),
+            make_dynamic_cache(
+                [
+                    (
+                        torch.randn(
+                            batch_size, encoder_attention_heads, max_source_positions, head_dim
+                        ),
+                        torch.randn(
+                            batch_size, encoder_attention_heads, max_source_positions, head_dim
+                        ),
+                    )
+                    for i in range(num_hidden_layers)
+                ]
+            ),
+        ),
+        # one these is selected based on the forward method signature
+        # encoder_last_hidden_state=torch.randn(batch_size, sequence_length2, encoder_dim),
+        # encoder_outputs=torch.randn(batch_size, sequence_length2, encoder_dim),
+    )
+    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,
+            dummy_max_token_id=dummy_max_token_id,
+            max_source_positions=max_source_positions,
+            d_model=d_model,
+            num_hidden_layers=num_hidden_layers,
+            encoder_attention_heads=encoder_attention_heads,
+            encoder_layers=encoder_layers,
+            decoder_layers=decoder_layers,
+            head_dim=head_dim,
+            batch_size=batch_size + 1,
+            sequence_length=sequence_length + add_second_input,
+            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,
+            "d_model",
+            "decoder_attention_heads",
+            "decoder_layers",
+            "encoder_attention_heads",
+            "encoder_layers",
+            "max_source_positions",
+            "num_hidden_layers",
+            "vocab_size",
+        )
+    kwargs = dict(
+        batch_size=2,
+        sequence_length=30,
+        dummy_max_token_id=31000 if config is None else config.vocab_size,
+        max_source_positions=1500 if config is None else config.max_source_positions,
+        d_model=384 if config is None else config.d_model,
+        num_hidden_layers=4 if config is None else config.num_hidden_layers,
+        encoder_attention_heads=6 if config is None else config.encoder_attention_heads,
+        encoder_layers=4 if config is None else config.encoder_layers,
+        decoder_attention_heads=6 if config is None else config.decoder_attention_heads,
+        decoder_layers=4 if config is None else config.decoder_layers,
+        head_dim=(
+            64 if config is None else (config.d_model // config.encoder_attention_heads)
+        ),
+    )
+    return kwargs, get_inputs

onnx_diagnostic/tasks/data/__init__.py

@@ -0,0 +1,13 @@
+import os
+
+
+def get_data(name: str):
+    """Returns data stored in this folder."""
+    filename = os.path.join(os.path.dirname(__file__), name)
+    assert os.path.exists(
+        filename
+    ), f"Unable to find a file with {name!r}, looked for {filename!r}"
+
+    from ...helpers.mini_onnx_builder import create_input_tensors_from_onnx_model
+
+    return create_input_tensors_from_onnx_model(filename)

onnx_diagnostic/tasks/feature_extraction.py

@@ -0,0 +1,162 @@
+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,
+)
+from ..helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
+
+
+__TASK__ = "feature-extraction"
+
+
+def reduce_model_config(config: Any) -> Dict[str, Any]:
+    """Reduces a model size."""
+    check_hasattr(config, "num_hidden_layers")
+    kwargs = dict(num_hidden_layers=min(config.num_hidden_layers, nhl()))
+    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,
+    sequence_length2: int = 3,
+    decoder_attention_heads: Optional[int] = None,
+    encoder_attention_heads: Optional[int] = None,
+    encoder_ffn_dim: Optional[int] = None,
+    decoder_ffn_dim: Optional[int] = None,
+    num_hidden_layers: Optional[int] = None,
+    add_second_input: int = 1,
+    **kwargs,  # unused
+):
+    """
+    Generates inputs for task ``feature-extraction``.
+    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},
+        "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
+        ),
+        attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64),
+    )
+    if (
+        encoder_attention_heads
+        and decoder_attention_heads
+        and encoder_ffn_dim
+        and decoder_ffn_dim
+        and num_hidden_layers
+    ):
+        inputs["past_key_values"] = make_encoder_decoder_cache(
+            make_dynamic_cache(
+                [
+                    (
+                        torch.randn(
+                            batch_size,
+                            encoder_attention_heads,
+                            sequence_length,
+                            encoder_ffn_dim,
+                        ),
+                        torch.randn(
+                            batch_size,
+                            encoder_attention_heads,
+                            sequence_length,
+                            encoder_ffn_dim,
+                        ),
+                    )
+                    for i in range(num_hidden_layers)
+                ]
+            ),
+            make_dynamic_cache(
+                [
+                    (
+                        torch.randn(
+                            batch_size,
+                            decoder_attention_heads,
+                            sequence_length2,
+                            decoder_ffn_dim,
+                        ),
+                        torch.randn(
+                            batch_size,
+                            decoder_attention_heads,
+                            sequence_length2,
+                            decoder_ffn_dim,
+                        ),
+                    )
+                    for i in range(num_hidden_layers)
+                ]
+            ),
+        )
+        cache_length = "cache_length_key"
+        cache_length2 = "cache_length_val"
+        shapes["past_key_values"] = [  # type: ignore[assignment]
+            [{0: batch, 2: cache_length} for _ in range(num_hidden_layers * 2)],
+            [{0: batch, 2: cache_length2} for _ in range(num_hidden_layers * 2)],
+        ]
+
+    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,
+            sequence_length2=sequence_length2,
+            decoder_attention_heads=decoder_attention_heads,
+            encoder_attention_heads=encoder_attention_heads,
+            encoder_ffn_dim=encoder_ffn_dim,
+            decoder_ffn_dim=decoder_ffn_dim,
+            num_hidden_layers=num_hidden_layers,
+            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),
+    )
+    for att in [
+        "decoder_attention_heads",
+        "encoder_attention_heads",
+        "encoder_ffn_dim",
+        "decoder_ffn_dim",
+        "num_hidden_layers",
+    ]:
+        if hasattr(config, att):
+            kwargs[att] = getattr(config, att)
+    kwargs["decoder_ffn_dim"] = kwargs["encoder_ffn_dim"] = 64
+    return kwargs, get_inputs