onnx-diagnostic 0.3.0__py3-none-any.whl → 0.4.0__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.3.0"
+__version__ = "0.4.0"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -1,3 +1,4 @@
+import argparse
 import json
 import sys
 import textwrap
@@ -227,6 +228,21 @@ def _cmd_config(argv: List[Any]):
         print(f"task: {task_from_id(args.mid)}")
 
 
+class _ParseDict(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        d = getattr(namespace, self.dest) or {}
+
+        if values:
+            for item in values:
+                split_items = item.split("=", 1)
+                key = split_items[0].strip()  # we remove blanks around keys, as is logical
+                value = split_items[1]
+
+                d[key] = value
+
+        setattr(namespace, self.dest, d)
+
+
 def get_parser_validate() -> ArgumentParser:
     parser = ArgumentParser(
         prog="test",
@@ -287,22 +303,37 @@ def get_parser_validate() -> ArgumentParser:
         help="drops the following inputs names, it should be a list "
         "with comma separated values",
     )
+    parser.add_argument(
+        "--ortfusiontype",
+        required=False,
+        help="applies onnxruntime fusion, this parameter should contain the "
+        "model type or multiple values separated by `|`. `ALL` can be used "
+        "to run them all",
+    )
     parser.add_argument("-v", "--verbose", default=0, type=int, help="verbosity")
     parser.add_argument("--dtype", help="changes dtype if necessary")
     parser.add_argument("--device", help="changes the device if necessary")
+    parser.add_argument(
+        "--iop",
+        metavar="KEY=VALUE",
+        nargs="*",
+        help="Additional input options, use to change the default "
+        "inputs use to export, example: --iop cls_cache=SlidingWindowCache",
+        action=_ParseDict,
+    )
     return parser
 
 
 def _cmd_validate(argv: List[Any]):
     from .helpers import string_type
-    from .torch_models.test_helper import get_inputs_for_task, validate_model, _ds_clean
-    from .torch_models.hghub.model_inputs import get_get_inputs_function_for_tasks
+    from .torch_models.test_helper import get_inputs_for_task, validate_model
+    from .tasks import supported_tasks
 
     parser = get_parser_validate()
     args = parser.parse_args(argv[1:])
     if not args.task and not args.mid:
         print("-- list of supported tasks:")
-        print("\n".join(sorted(get_get_inputs_function_for_tasks())))
+        print("\n".join(supported_tasks()))
     elif not args.mid:
         data = get_inputs_for_task(args.task)
         if args.verbose:
@@ -313,8 +344,16 @@ def _cmd_validate(argv: List[Any]):
             print(f"  + {k.ljust(max_length)}: {string_type(v, with_shape=True)}")
         print("-- dynamic_shapes")
         for k, v in data["dynamic_shapes"].items():
-            print(f"  + {k.ljust(max_length)}: {_ds_clean(v)}")
+            print(f"  + {k.ljust(max_length)}: {string_type(v)}")
     else:
+        # Let's skip any invalid combination if known to be unsupported
+        if (
+            "onnx" not in (args.export or "")
+            and "custom" not in (args.export or "")
+            and (args.opt or "")
+        ):
+            print(f"validate - unsupported args: export={args.export!r}, opt={args.opt!r}")
+            return
         summary, _data = validate_model(
             model_id=args.mid,
             task=args.task,
@@ -330,6 +369,8 @@ def _cmd_validate(argv: List[Any]):
             exporter=args.export,
             dump_folder=args.dump_folder,
             drop_inputs=None if not args.drop else args.drop.split(","),
+            ortfusiontype=args.ortfusiontype,
+            input_options=args.iop,
         )
         print("")
         print("-- summary --")

onnx_diagnostic/export/__init__.py

@@ -1 +1,2 @@
 from .dynamic_shapes import CoupleInputsDynamicShapes, ModelInputs
+from .validate import validate_ep

onnx_diagnostic/export/dynamic_shapes.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
 import torch
 from ..helpers import string_type
@@ -8,6 +8,30 @@ from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
 DYNAMIC_SHAPES = Tuple[Tuple[Any, ...], Dict[str, Any]]
 
 
+def flatten_dynamic_shapes(ds: Any) -> Any:
+    """Flattens the dynamic shapes."""
+    if isinstance(ds, list):
+        return _flat_list([flatten_dynamic_shapes(t) for t in ds])
+    if isinstance(ds, tuple):
+        return tuple(_flat_list([flatten_dynamic_shapes(t) for t in ds]))
+    if isinstance(ds, dict):
+        if all(isinstance(i, int) for i in ds):
+            # That's a dynamic shape
+            return ds
+        return _flat_list([flatten_dynamic_shapes(t) for t in ds.values()])
+    raise AssertionError(f"Not implemented for {type(ds)}: {ds}")
+
+
+def _flat_list(li: List[Any]) -> List[Dict[int, str]]:
+    res = []
+    for t in li:
+        if isinstance(t, dict):
+            res.append(t)
+        else:
+            res.extend(t)
+    return res
+
+
 class CoupleInputsDynamicShapes:
     """
     Pair inputs / dynamic shapes.
@@ -68,7 +92,8 @@ class CoupleInputsDynamicShapes:
         return self._generic_walker(
             lambda inputs, ds, value=value: self._replace_string_dim_tensor(
                 inputs, ds, value=value
-            )
+            ),
+            flatten_unflatten=True,
         )
 
     @classmethod
@@ -76,7 +101,7 @@ class CoupleInputsDynamicShapes:
         assert isinstance(inputs, torch.Tensor), f"unexpected type for inputs {type(inputs)}"
         assert isinstance(ds, dict) and all(isinstance(s, int) for s in ds), (
             f"Unexpected types, inputs is a Tensor but ds is {ds}, "
-            f"a dictionary is expected to specify a dimension dimension"
+            f"a dictionary is expected to specify a dimension"
         )
         if value is None:
             value = torch.export.Dim.DYNAMIC
@@ -86,6 +111,57 @@ class CoupleInputsDynamicShapes:
                 new_ds[i] = value
         return new_ds
 
+    def replace_by_string(self):
+        """
+        Replaces dimensions by strings.
+
+        Example:
+
+        .. runpython::
+            :showcode:
+
+            import torch
+            from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes
+
+            Dim = torch.export.Dim
+            T3x1 = torch.rand((3, 1))
+            T3x4 = torch.rand((3, 4))
+            ds_batch = {0: Dim("batch")}
+            ds_batch_seq = {0: Dim("batch"), 1: Dim("seq")}
+            kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
+            ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
+            print(CoupleInputsDynamicShapes((), kwargs, ds).replace_by_string())
+        """
+        unique = set()
+        return self._generic_walker(
+            lambda inputs, ds, unique=unique: self._replace_dim_tensor_by_string(
+                inputs, ds, unique=unique
+            ),
+            flatten_unflatten=True,
+        )
+
+    @classmethod
+    def _replace_dim_tensor_by_string(cls, inputs, ds, unique: Set[str]):
+        assert isinstance(inputs, torch.Tensor), f"unexpected type for inputs {type(inputs)}"
+        assert isinstance(ds, dict) and all(isinstance(s, int) for s in ds), (
+            f"Unexpected types, inputs is a Tensor but ds is {ds}, "
+            f"a dictionary is expected to specify a dimension"
+        )
+        new_ds = ds.copy()
+        for i, v in ds.items():
+            if isinstance(v, str):
+                unique.add(v)
+                new_ds[i] = v
+            elif v in (torch.export.Dim.DYNAMIC, torch.export.Dim.AUTO):
+                name = f"Dim{len(unique)}"
+                new_ds[i] = name
+                unique.add(name)
+            else:
+                name = v.__name__
+                unique.add(name)
+                new_ds[i] = name
+        return new_ds
+
     def invalid_dimensions_for_export(self):
         """
         Tells if the inputs are valid based on the dynamic shapes definition.
@@ -129,7 +205,7 @@ class CoupleInputsDynamicShapes:
             ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
             print(CoupleInputsDynamicShapes((), kwargs, ds).invalid_dimensions_for_export())
         """
-        return self._generic_walker(self._valid_shapes_tensor)
+        return self._generic_walker(self._valid_shapes_tensor, flatten_unflatten=True)
 
     @classmethod
     def _valid_shapes_tensor(cls, inputs, ds):
@@ -147,7 +223,9 @@ class CoupleInputsDynamicShapes:
                     issues[i] = f"d=[{d}]"
         return issues if issues else None
 
-    def _generic_walker(self, processor: Callable):
+    def _generic_walker(
+        self, processor: Callable, args_kwargs: bool = False, flatten_unflatten: bool = False
+    ):
         """
         Generic deserializator walking through inputs and dynamic_shapes all along.
         The function returns a result with the same structure as the dynamic shapes.
@@ -157,14 +235,23 @@ class CoupleInputsDynamicShapes:
                 f"Type mismatch, args={string_type(self.args)} and "
                 f"dynamic_shapes={self.dynamic_shapes} should have the same type."
             )
-            return self._generic_walker_step(processor, self.kwargs, self.dynamic_shapes)
+            res = self._generic_walker_step(
+                processor,
+                self.kwargs,
+                self.dynamic_shapes,
+                flatten_unflatten=flatten_unflatten,
+            )
+            return (tuple(), res) if args_kwargs else res
 
         if not self.kwargs:
             assert isinstance(self.args, tuple) and isinstance(self.dynamic_shapes, tuple), (
                 f"Type mismatch, args={string_type(self.args)} and "
                 f"dynamic_shapes={self.dynamic_shapes} should have the same type."
             )
-            return self._generic_walker_step(processor, self.args, self.dynamic_shapes)
+            res = self._generic_walker_step(
+                processor, self.args, self.dynamic_shapes, flatten_unflatten=flatten_unflatten
+            )
+            return (res, {}) if args_kwargs else res
 
         assert isinstance(self.dynamic_shapes, dict), (
             f"Both positional and named arguments (args and kwargs) are filled. "
@@ -174,12 +261,22 @@ class CoupleInputsDynamicShapes:
             self.dynamic_shapes
         ):
             # No dynamic shapes for the positional arguments.
-            return self._generic_walker_step(processor, self.kwargs, self.dynamic_shapes)
+            return self._generic_walker_step(
+                processor,
+                self.kwargs,
+                self.dynamic_shapes,
+                flatten_unflatten=flatten_unflatten,
+            )
 
         if isinstance(self.args_names, list):
             if not set(self.args_names) & set(self.dynamic_shapes):
                 # No dynamic shapes for the positional arguments.
-                return self._generic_walker_step(processor, self.kwargs, self.dynamic_shapes)
+                return self._generic_walker_step(
+                    processor,
+                    self.kwargs,
+                    self.dynamic_shapes,
+                    flatten_unflatten=flatten_unflatten,
+                )
 
             assert self.args_names, (
                 "args and kwargs are filled, then args_names must be specified in "
@@ -192,7 +289,19 @@ class CoupleInputsDynamicShapes:
             )
             kwargs = dict(zip(self.args_names, self.args))
             kwargs.update(self.kwargs)
-            return self._generic_walker_step(processor, kwargs, self.dynamic_shapes)
+            res = self._generic_walker_step(
+                processor, kwargs, self.dynamic_shapes, flatten_unflatten=flatten_unflatten
+            )
+            if args_kwargs:
+                pgs = [None for _ in range(len(self.args))]
+                kws = {}
+                for k, v in res.items():
+                    if k not in self.kwargs:
+                        pgs[self.args_names.index(k)] = v
+                    else:
+                        kws[k] = v
+                return pgs, kws
+            return res
 
         raise NotImplementedError(
             f"Not yet implemented when args is filled, "
@@ -200,35 +309,48 @@ class CoupleInputsDynamicShapes:
         )
 
     @classmethod
-    def _generic_walker_step(cls, processor: Callable, inputs, ds):
+    def _generic_walker_step(
+        cls, processor: Callable, inputs, ds, flatten_unflatten: bool = False
+    ):
         if isinstance(inputs, torch.Tensor):
             return processor(inputs, ds)
         if isinstance(inputs, (int, float, str)):
             return None
-        if isinstance(inputs, (tuple, list, dict)):
-            assert type(ds) is type(
-                inputs
-            ), f"Type mismatch between inputs {type(inputs)} and ds={type(ds)}"
+        if type(inputs) in (tuple, list, dict):
+            # Type must be strict, some custom classes can inherit from those.
+            assert type(inputs) is type(ds), (
+                f"Input type and dynamic shape type mush match but "
+                f"type(inputs)={type(inputs)}, type(ds)={type(ds)}, "
+                f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
+            )
             assert len(ds) == len(inputs), (
                 f"Length mismatch between inputs {len(inputs)} "
                 f"and ds={len(ds)}\n"
                 f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
             )
-            if isinstance(inputs, (tuple, list)):
+            if type(inputs) in (tuple, list):
                 value = []
                 for i, d in zip(inputs, ds):
-                    value.append(cls._generic_walker_step(processor, i, d))
+                    value.append(
+                        cls._generic_walker_step(
+                            processor, i, d, flatten_unflatten=flatten_unflatten
+                        )
+                    )
                 return (
                     (value if isinstance(ds, list) else tuple(value))
                     if any(v is not None for v in value)
                     else None
                 )
-            assert set(inputs) == set(
-                ds
-            ), f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}"
+            assert type(inputs) is dict, f"Unexpected type for inputs {type(inputs)}"
+            assert set(inputs) == set(ds), (
+                f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}, "
+                f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
+            )
             dvalue = {}
             for k, v in inputs.items():
-                t = cls._generic_walker_step(processor, v, ds[k])
+                t = cls._generic_walker_step(
+                    processor, v, ds[k], flatten_unflatten=flatten_unflatten
+                )
                 if t is not None:
                     dvalue[k] = t
             return dvalue if dvalue else None
@@ -239,12 +361,22 @@ class CoupleInputsDynamicShapes:
             f"torch.utils._pytree.register_pytree_node, it is not possible to "
             f"map this class with the given dynamic shapes."
         )
+        if flatten_unflatten:
+            flatunflat = flatten_unflatten_for_dynamic_shapes(inputs)
+            return cls._generic_walker_step(
+                processor, flatunflat, ds, flatten_unflatten=flatten_unflatten
+            )
         flat, _spec = torch.utils._pytree.tree_flatten(inputs)
-        return cls._generic_walker_step(processor, flat, ds)
+        if all(isinstance(t, torch.Tensor) for t in flat):
+            # We need to flatten dynamic shapes as well
+            ds = flatten_dynamic_shapes(ds)
+        return cls._generic_walker_step(
+            processor, flat, ds, flatten_unflatten=flatten_unflatten
+        )
 
     class ChangeDimensionProcessor:
-        def __init__(self):
-            self.mapping = {}
+        def __init__(self, desired_values):
+            self.mapping = desired_values or {}
 
         def _build_new_shape(
             self, shape: Tuple[int, ...], ds: Dict[int, Any]
@@ -285,14 +417,14 @@ class CoupleInputsDynamicShapes:
                     tuple(alt_shape), dtype=tensor.dtype, device=tensor.device
                 )
                 mind = min(d0, d1)
-                indices = [slice(None) for _ in range(rank)]
+                indices: List[Union[slice, int]] = [slice(None) for _ in range(rank)]
                 indices[i] = slice(0, mind)
                 ind = tuple(indices)
                 new_tensor[ind] = tensor[ind]
                 if d1 > mind:
                     for k in range(d1 - mind):
-                        indices0 = [slice(None) for _ in range(rank)]
-                        indices1 = [slice(None) for _ in range(rank)]
+                        indices0: List[Union[slice, int]] = [slice(None) for _ in range(rank)]
+                        indices1: List[Union[slice, int]] = [slice(None) for _ in range(rank)]
                         indices1[i] = mind + k
                         indices0[i] = k % mind
                         new_tensor[tuple(indices1)] = tensor[tuple(indices0)]
@@ -310,7 +442,9 @@ class CoupleInputsDynamicShapes:
             new_shape = self._build_new_shape(inputs.shape, ds)
             return self._build_new_tensor(inputs, new_shape)
 
-    def change_dynamic_dimensions(self):
+    def change_dynamic_dimensions(
+        self, desired_values: Optional[Dict[str, int]] = None, args_kwargs: bool = False
+    ):
         """
         A model exported with dynamic shapes is not necessarily dynamic
         just because the user specified dynamic shapes. The algorithm
@@ -320,6 +454,10 @@ class CoupleInputsDynamicShapes:
         for the dimension than the first ones, assuming they were used to export
         the model.
 
+        :param desired_values: to fixed named dimension to have the desired value
+        :param args_kwargs: return both args, kwargs even if empty
+        :return: new inputs
+
         Example:
 
         .. runpython::
@@ -340,7 +478,9 @@ class CoupleInputsDynamicShapes:
             print("before:", string_type(kwargs, with_shape=True))
             print("-after:", string_type(new_kwargs, with_shape=True))
         """
-        return self._generic_walker(self.ChangeDimensionProcessor())
+        return self._generic_walker(
+            self.ChangeDimensionProcessor(desired_values), args_kwargs=args_kwargs
+        )
 
 
 class ModelInputs:

onnx_diagnostic/export/validate.py

@@ -0,0 +1,170 @@
+import inspect
+import itertools
+import time
+from typing import Any, Dict, List, Optional, Tuple, Union
+import torch
+from ..helpers import string_type, max_diff, string_diff
+from ..helpers.torch_test_helper import torch_deepcopy
+from .dynamic_shapes import CoupleInputsDynamicShapes
+
+
+def compare_modules(
+    modep: torch.nn.Module,
+    mod: Optional[torch.nn.Module] = None,
+    args: Optional[Tuple[Any, ...]] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    copy: bool = False,
+    exc: bool = True,
+    verbose: int = 0,
+    atol: float = 1e-2,
+    rtol: float = 1e-1,
+) -> Dict[str, Any]:
+    """
+    Compares two torch modules, usually one coming from an exported program,
+    the other being the origin model.
+
+    :param model: first module
+    :param mod: second module (it produces the expected values)
+    :param args: positional arguments
+    :param kwargs: named arguments
+    :param copy: copy the inputs before executing the model (they may modify them inplace)
+    :param exc: raise exception if discrepancies are too high
+    :param verbose: verbosity level
+    :param atol: absolute tolerance
+    :param rtol: relative tolerance
+    :return: dictionary with inputs, outputs and tolerance
+
+    Example:
+
+    .. runpython::
+        :showcode:
+
+        import torch
+        from onnx_diagnostic.export import validate_ep, CoupleInputsDynamicShapes
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y
+
+        model = Model()
+        x = torch.randn((5, 6))
+        y = torch.randn((1, 6))
+        model(x, y)  # to make it is running
+
+        ds = ({0: "a", 1: "b"}, {1: "b"})
+        cpl = CoupleInputsDynamicShapes((x, y), {}, ds)
+        ep = torch.export.export(model, (x, y), dynamic_shapes=cpl.replace_string_by())
+        validate_ep(
+            ep,
+            model,
+            args=(x, y),
+            verbose=2,
+            copy=True,
+            dynamic_shapes=ds,
+            values_to_try={"a": [5, 10], "b": [10, 20]},
+        )
+
+    """
+    args = args or ()
+    kwargs = kwargs or {}
+
+    def _get(a):
+        return torch_deepcopy(a) if copy else a
+
+    if verbose:
+        begin = time.perf_counter()
+        print(
+            f"[compare_modules] check ep with "
+            f"args={string_type(args, with_shape=True, with_device=True)}, "
+            f"kwargs={string_type(kwargs, with_shape=True, with_device=True)}..."
+        )
+    got = modep(*_get(args), **_get(kwargs))
+    if verbose:
+        d = time.perf_counter() - begin
+        print(f"[compare_modules] done in {d} with output={string_type(got, with_shape=True)}")
+    if mod:
+        if verbose:
+            begin = time.perf_counter()
+            print("[compare_modules] run torch module...")
+        expected = mod(*_get(args), **_get(kwargs))
+        diff = max_diff(expected, got)
+        if verbose:
+            d = time.perf_counter() - begin
+            print(
+                f"[compare_modules] done in {d} with "
+                f"output={string_type(expected, with_shape=True)}"
+            )
+            print(f"[compare_modules] discrepancies={string_diff(diff)}")
+        assert not exc or (
+            diff["abs"] <= atol and diff["rel"] <= rtol
+        ), f"Discrepancies={string_diff(diff)} higher than expected."
+        return dict(args=args, kwargs=kwargs, expected=expected, got=got, diff=diff)
+    return dict(args=args, kwargs=kwargs, got=got)
+
+
+def validate_ep(
+    ep: Union[torch.nn.Module, torch.export.ExportedProgram],
+    mod: Optional[torch.nn.Module] = None,
+    args: Optional[Tuple[Any, ...]] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
+    copy: bool = False,
+    dynamic_shapes: Optional[Any] = None,
+    values_to_try: Optional[Dict[str, List[int]]] = None,
+    exc: bool = True,
+    verbose: int = 0,
+    atol: float = 1e-2,
+    rtol: float = 1e-1,
+) -> List[Dict[str, Any]]:
+    """
+    Validates an exported program.
+
+    :param model: first module
+    :param mod: second module (it produces the expected values)
+    :param args: positional arguments
+    :param kwargs: named arguments
+    :param copy: copy the inputs before executing the model (they may modify them inplace)
+    :param dynamic_shapes: dynamic shapes, string should be used not ``torch.export.Dim``
+    :param values_to_try: dictionary with the values to try for every dynamic dimension
+    :param exc: raise exception if discrepancies are too high
+    :param verbose: verbosity level
+    :param atol: absolute tolerance
+    :param rtol: relative tolerance
+    :return: dictionary with inputs, outputs and tolerance
+    """
+    modep = ep.module() if isinstance(ep, torch.export.ExportedProgram) else ep
+
+    results = [
+        compare_modules(
+            modep, mod, args, kwargs, copy=copy, verbose=verbose, atol=atol, rtol=rtol
+        )
+    ]
+
+    assert (dynamic_shapes and values_to_try) or (
+        not dynamic_shapes and not values_to_try
+    ), "Either both dynamic_shapes and values_to_try are specified, either none."
+    if not dynamic_shapes or not values_to_try:
+        return results
+
+    items = list(values_to_try.items())
+    keys = [_[0] for _ in items]
+    values = [_[1] for _ in items]
+    all_vals = list(itertools.product(*values))
+    cpl = CoupleInputsDynamicShapes(
+        args or (),
+        kwargs or {},
+        dynamic_shapes,
+        args_names=(
+            list(inspect.signature(modep.forward).parameters) if args and kwargs else None
+        ),
+    )
+    for i, vals in enumerate(all_vals):
+        change_dims = dict(zip(keys, vals))
+        if verbose:
+            print(f"[validate_ep] try {i}/{len(all_vals)}: {change_dims}")
+        new_params = cpl.change_dynamic_dimensions(change_dims, args_kwargs=True)
+        na, nkw = new_params
+        c = compare_modules(
+            modep, mod, na, nkw, copy=copy, verbose=max(verbose - 1, 0), atol=atol, rtol=rtol
+        )
+        results.append(c)
+    return results