onnx-diagnostic 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

onnx_diagnostic/torch_models/test_helper.py

@@ -4,6 +4,7 @@ import os
 import sys
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import time
+import numpy as np
 import onnx
 import onnxscript
 import onnxscript.rewriter.ort_fusions as ort_fusions
@@ -17,6 +18,7 @@ from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
 from ..tasks import random_input_kwargs
 from ..torch_export_patches import torch_export_patches
 from ..torch_export_patches.patch_inputs import use_dyn_not_str
+from ..reference import TorchOnnxEvaluator
 from .hghub import get_untrained_model_with_inputs
 
 
@@ -192,11 +194,16 @@ def _quiet_or_not_quiet(
     summary: Dict[str, Any],
     data: Optional[Dict[str, Any]],
     fct: Callable,
+    repeat: int = 1,
+    warmup: int = 0,
 ) -> Any:
     begin = time.perf_counter()
     if quiet:
         try:
-            return fct()
+            res = fct()
+            summary[f"time_{suffix}"] = time.perf_counter() - begin
+            if warmup + repeat == 1:
+                return res
         except Exception as e:
             summary[f"ERR_{suffix}"] = str(e)
             summary[f"time_{suffix}"] = time.perf_counter() - begin
@@ -204,11 +211,45 @@ def _quiet_or_not_quiet(
                 return {f"ERR_{suffix}": e}
             data[f"ERR_{suffix}"] = e
             return None
-    res = fct()
+    else:
+        res = fct()
     summary[f"time_{suffix}"] = time.perf_counter() - begin
+    if warmup + repeat > 1:
+        if suffix == "run":
+            res = torch_deepcopy(res)
+        summary[f"{suffix}_output"] = string_type(res, with_shape=True, with_min_max=True)
+        summary[f"{suffix}_warmup"] = warmup
+        summary[f"{suffix}_repeat"] = repeat
+        for _w in range(max(0, warmup - 1)):
+            t = fct()
+            summary[f"io_{suffix}_{_w+1}"] = string_type(t, with_shape=True, with_min_max=True)
+        summary[f"time_{suffix}_warmup"] = time.perf_counter() - begin
+        times = []
+        for _r in range(repeat):
+            begin = time.perf_counter()
+            t = fct()
+            times.append(time.perf_counter() - begin)
+        a = np.array(times)
+        summary[f"time_{suffix}_latency"] = a.mean()
+        summary[f"time_{suffix}_latency_std"] = a.std()
+        summary[f"time_{suffix}_latency_min"] = a.min()
+        summary[f"time_{suffix}_latency_min"] = a.max()
     return res
 
 
+def shrink_config(cfg: Dict[str, Any]) -> Dict[str, Any]:
+    """Shrinks the configuration before it gets added to the information to log."""
+    new_cfg = {}
+    for k, v in cfg.items():
+
+        new_cfg[k] = (
+            v
+            if (not isinstance(v, (list, tuple, set, dict)) or len(v) < 50)
+            else (v.__class__("...") if isinstance(v, (list, tuple)) else "...")
+        )
+    return new_cfg
+
+
 def validate_model(
     model_id: str,
     task: Optional[str] = None,
@@ -231,9 +272,14 @@ def validate_model(
     model_options: Optional[Dict[str, Any]] = None,
     subfolder: Optional[str] = None,
     opset: Optional[int] = None,
+    runtime: str = "onnxruntime",
+    repeat: int = 1,
+    warmup: int = 0,
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Validates a model.
+    The function can also be called through the command line
+    :ref:`l-cmd-validate`.
 
     :param model_id: model id to validate
     :param task: task used to generate the necessary inputs,
@@ -241,7 +287,8 @@ def validate_model(
         if it can be determined
     :param do_run: checks the model works with the defined inputs
     :param exporter: exporter the model using this exporter,
-        available list: ``export-strict``, ``export-nostrict``, ``onnx``
+        available list: ``export-strict``, ``export-nostrict``, ...
+        see below
     :param do_same: checks the discrepancies of the exported model
     :param verbose: verbosity level
     :param dtype: uses this dtype to check the model
@@ -267,6 +314,10 @@ def validate_model(
         ``num_hidden_layers`` or ``attn_implementation``
     :param subfolder: version or subfolders to uses when retrieving a model id
     :param opset: onnx opset to use for the conversion
+    :param runtime: onnx runtime to use to check about discrepancies,
+        only if `do_run` is true
+    :param repeat: number of time to measure the model
+    :param warmup: warmup the model first
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
 
@@ -274,6 +325,20 @@ def validate_model(
     information:
 
     * ``PRINT_CONFIG``: prints the model configuration
+
+    The following exporters are available:
+
+    * ``export-nostrict``: run :func:`torch.export.export` (..., strict=False)
+    * ``onnx-dynamo``: run :func:`torch.onnx.export` (..., dynamo=True),
+      models can be optimized with ``optimization`` in ``("ir", "os_ort")``
+    * ``modelbuilder``: use :epkg:`ModelBuilder` to builds the onnx model
+    * ``custom``: custom exporter (see :epkg:`experimental-experiment`),
+      models can be optimized with ``optimization`` in
+      ``("default", "default+onnxruntime", "default+os_ort", "default+onnxruntime+os_ort")``
+
+    The default runtime, :epkg:`onnxruntime` is used to validate a model and check the
+    exported model returns the same outputs as the original one, otherwise,
+    :class:`onnx_diagnostic.reference.TorchOnnxEvaluator` is used.
     """
     assert (
         not rewrite or patch
@@ -295,6 +360,7 @@ def validate_model(
             version_ortfusiontype=ortfusiontype or "",
             version_stop_if_static=str(stop_if_static),
             version_exporter=exporter or "",
+            version_runtime=runtime,
         )
     )
     if opset:
@@ -436,7 +502,9 @@ def validate_model(
     if summary["model_module"] in sys.modules:
         summary["model_file"] = str(sys.modules[summary["model_module"]].__file__)  # type: ignore[index]
     summary["model_config_class"] = data["configuration"].__class__.__name__
-    summary["model_config"] = str(data["configuration"].to_dict()).replace(" ", "")
+    summary["model_config"] = str(shrink_config(data["configuration"].to_dict())).replace(
+        " ", ""
+    )
     summary["model_id"] = model_id
 
     if verbose:
@@ -460,7 +528,13 @@ def validate_model(
         model = data["model"]
 
         expected = _quiet_or_not_quiet(
-            quiet, "run", summary, data, (lambda m=model, inp=inputs: m(**inp))
+            quiet,
+            "run",
+            summary,
+            data,
+            (lambda m=model, inp=inputs: m(**torch_deepcopy(inp))),
+            repeat=repeat,
+            warmup=warmup,
         )
         if "ERR_run" in summary:
             return summary, data
@@ -522,7 +596,7 @@ def validate_model(
 
                     disc = max_diff(data["expected"], expected)
                     for k, v in disc.items():
-                        summary[f"disc_patched_{k}"] = v
+                        summary[f"disc_patched_{k}"] = str(v)
                     if verbose:
                         print("[validate_model] done (patched run)")
                         print(f"[validate_model] patched discrepancies={string_diff(disc)}")
@@ -618,7 +692,14 @@ def validate_model(
         return summary, data
 
     if do_run:
-        summary_valid, data = validate_onnx_model(data=data, quiet=quiet, verbose=verbose)
+        summary_valid, data = validate_onnx_model(
+            data=data,
+            quiet=quiet,
+            verbose=verbose,
+            runtime=runtime,
+            repeat=repeat,
+            warmup=warmup,
+        )
         summary.update(summary_valid)
 
     if ortfusiontype and "onnx_filename" in data:
@@ -671,7 +752,13 @@ def validate_model(
 
             if do_run:
                 summary_valid, data = validate_onnx_model(
-                    data=data, quiet=quiet, verbose=verbose, flavour=flavour
+                    data=data,
+                    quiet=quiet,
+                    verbose=verbose,
+                    flavour=flavour,
+                    runtime=runtime,
+                    repeat=repeat,
+                    warmup=warmup,
                 )
                 summary.update(summary_valid)
 
@@ -883,6 +970,9 @@ def validate_onnx_model(
     quiet: bool = False,
     verbose: int = 0,
     flavour: Optional[str] = None,
+    runtime: str = "onnxruntime",
+    repeat: int = 1,
+    warmup: int = 0,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """
     Verifies that an onnx model produces the same
@@ -895,6 +985,9 @@ def validate_onnx_model(
     :param quiet: catch exception or not
     :param verbose: verbosity
     :param flavour: use a different version of the inputs
+    :param runtime: onnx runtime to use, onnxruntime or torch
+    :param repeat: run that number of times the model
+    :param warmup: warmup the model
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -936,18 +1029,28 @@ def validate_onnx_model(
             f"{providers}..., flavour={flavour!r}"
         )
 
+    cls_runtime = (
+        (
+            lambda model, providers: onnxruntime.InferenceSession(
+                (model.SerializeToString() if isinstance(model, onnx.ModelProto) else model),
+                providers=providers,
+            )
+        )
+        if runtime == "onnxruntime"
+        else (
+            lambda model, providers: TorchOnnxEvaluator(
+                model, providers=providers, verbose=max(verbose - 1, 0)
+            )
+        )
+    )
     sess = _quiet_or_not_quiet(
         quiet,
-        _mk("time_onnx_ort_create"),
+        _mk("onnx_ort_create"),
         summary,
         data,
-        (
-            lambda source=source, providers=providers: onnxruntime.InferenceSession(
-                source, providers=providers
-            )
-        ),
+        (lambda source=source, providers=providers: cls_runtime(source, providers)),
     )
-    if f"ERR_{_mk('time_onnx_ort_create')}" in summary:
+    if f"ERR_{_mk('onnx_ort_create')}" in summary:
         return summary, data
 
     data[_mk("onnx_ort_sess")] = sess
@@ -975,6 +1078,8 @@ def validate_onnx_model(
         summary,
         data,
         (lambda sess=sess, feeds=feeds: sess.run(None, feeds)),
+        repeat=repeat,
+        warmup=warmup,
     )
     if f"ERR_{_mk('time_onnx_ort_run')}" in summary:
         return summary, data
@@ -1051,7 +1156,7 @@ def call_torch_export_onnx(
             dynamo=False,
             dynamic_axes={
                 k: v
-                for k, v in CoupleInputsDynamicShapes(args, kwargs, ds)
+                for k, v in CoupleInputsDynamicShapes(args, kwargs, ds)  # type: ignore[arg-type]
                 .replace_by_string()
                 .items()
                 if isinstance(v, dict)
@@ -1229,6 +1334,13 @@ def call_torch_export_custom(
         "custom-nostrict",
         "custom-nostrict-default",
         "custom-nostrict-all",
+        "custom-inline",
+        "custom-strict-inline",
+        "custom-strict-default-inline",
+        "custom-strict-all-inline",
+        "custom-nostrict-inline",
+        "custom-nostrict-default-inline",
+        "custom-nostrict-all-inline",
     }
     assert exporter in available, f"Unexpected value for exporter={exporter!r} in {available}"
     assert "model" in data, f"model is missing from data: {sorted(data)}"
@@ -1269,6 +1381,10 @@ def call_torch_export_custom(
         ),
         save_ep=(os.path.join(dump_folder, f"{exporter}.ep") if dump_folder else None),
     )
+    inline = "-inline" in exporter
+    if inline:
+        export_options.aten_as_function = set()
+
     options = OptimizationOptions(patterns=optimization) if optimization else None
     model = data["model"]
     kws = dict(
@@ -1279,6 +1395,7 @@ def call_torch_export_custom(
         large_model=True,
         return_optimize_report=True,
         verbose=max(verbose - 2, 0),
+        inline=inline,
     )
     if opset:
         kws["target_opset"] = opset

onnx_diagnostic/torch_onnx/runtime_info.py

@@ -0,0 +1,289 @@
+import enum
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+import onnx
+import torch
+from ..api import TensorLike
+from ..helpers import string_type
+
+
+class RuntimeValueKind(enum.IntEnum):
+    "Kind of result."
+
+    RESULT = 1
+    INITIALIZER = 3
+    INPUT = 5
+    OUTPUT = 9
+
+    def to_str(self) -> str:
+        for k, v in self.__class__.__dict__.items():
+            if v == int(self):
+                return k
+        raise RuntimeError(f"Unable to display {self!r}")
+
+
+class RuntimeDevice(enum.IntEnum):
+    "Device definition"
+
+    UNKNOWN = 0
+    NEW = 1
+    CPU = 2
+    CUDA = 4
+
+    def to_str(self) -> str:
+        for k, v in self.__class__.__dict__.items():
+            if v == int(self):
+                return k
+        raise RuntimeError(f"Unable to display {self!r}")
+
+
+class RuntimeValue:
+    """Describes a value used during the execution of a model."""
+
+    def __init__(
+        self,
+        name: str,
+        dtype: Optional[Any] = None,
+        shape: Optional[Tuple[Union[str, int], ...]] = None,
+        value: Optional[Any] = None,
+        first_used: Optional[int] = None,
+        last_used: Optional[int] = None,
+        created: Optional[int] = None,
+        is_shape: Optional[bool] = None,
+        kind: Optional[RuntimeValueKind] = None,
+        device: Optional[RuntimeDevice] = None,
+    ):
+        self.name = name
+        self.dtype = dtype
+        self.shape = shape
+        self.value = value
+        self.first_used = first_used
+        self.last_used = last_used
+        self.created = created
+        self.is_shape = is_shape
+        self.kind = kind
+        self.device = device
+
+    def __repr__(self) -> str:
+        "usual"
+        ad = {}
+        for att in [
+            "name",
+            "dtype",
+            "shape",
+            "first_used",
+            "last_used",
+            "is_shape",
+            "kind",
+            "created",
+            "device",
+        ]:
+            v = getattr(self, att)
+            if v is not None:
+                ad[att] = v
+        if self.value is not None:
+            ad["value"] = (
+                self.value.string_type()
+                if hasattr(self.value, "string_type")
+                else string_type(self.value, with_shape=True)
+            )
+        msg = ", ".join(
+            f"{name}={t.to_str()}" if hasattr(t, "to_str") else f"{name}={t}"
+            for name, t in ad.items()
+        )
+        return f"{self.__class__.__name__}({msg})"
+
+    @property
+    def has_value(self) -> bool:
+        "Tells if value is specified."
+        return self.value is not None
+
+    def string_type(self) -> str:
+        "Returns a string describing the value."
+        rows = []
+        if self.shape is not None:
+            rows.append(f"shape={self.shape}")
+        if self.is_shape is not None:
+            rows.append(f"is_shape={self.is_shape}")
+        if self.device is not None:
+            rows.append(f"device={self.device}")
+        text = f", {', '.join(rows)}" if rows else ""
+        if self.value is None:
+            return (
+                f"RuntimeValue(name={self.name!r}{text}"
+                f", dtype={self.dtype}, kind={self.kind})"
+            )
+        return (
+            f"RuntimeValue(name={self.name!r}, "
+            f"kind={self.kind}{text}, value={self.value.string_type()})"
+        )
+
+    def set_value(self, value: Union[torch.Tensor, TensorLike]):
+        """Sets the value."""
+        assert value is not None, "Use clean_value to set a value to None"
+        self.value = value
+        is_sequence = hasattr(value, "is_sequence") and value.is_sequence()
+        if self.dtype:
+            assert value is None or self.dtype == value.dtype, (
+                f"Unexpected dtype={value.dtype}, previous dtype was {self.dtype}, "
+                f"is_sequence={is_sequence}"
+            )
+        else:
+            self.dtype = value.dtype
+        self.shape = None if is_sequence else tuple(map(int, value.shape))
+
+    def clean_value(self):
+        """Sets value to None."""
+        self.value = None
+
+    @property
+    def is_output(self) -> bool:
+        "Tells if it is an output."
+        return self.kind == RuntimeValueKind.OUTPUT
+
+    @property
+    def is_input(self) -> bool:
+        "Tells if it is an input."
+        return self.kind == RuntimeValueKind.INPUT
+
+    @property
+    def is_initializer(self) -> bool:
+        "Tells if it is an initializer."
+        return self.kind == RuntimeValueKind.INITIALIZER
+
+
+def get_hidden_inputs(graph: onnx.GraphProto) -> Set[str]:
+    """
+    Returns the hidden inputs (inputs coming from an upper context)
+    used by a subgraph.
+    """
+    hidden = set()
+    memo = (
+        set(i.name for i in graph.initializer)
+        | set(i.name for i in graph.sparse_initializer)
+        | set(i.name for i in graph.input)
+    )
+    for node in graph.node:
+        for i in node.input:
+            if i not in memo:
+                hidden.add(i)
+        for att in node.attribute:
+            if att.type == onnx.AttributeProto.GRAPH and att.g:
+                hid = get_hidden_inputs(att.g)
+                less = set(h for h in hid if h not in memo)
+                hidden |= less
+        memo |= set(node.output)
+    return hidden
+
+
+def set_is_shape(
+    node: onnx.NodeProto, values: Dict[str, RuntimeValue], drop: Optional[Set[str]] = None
+) -> List[str]:
+    """
+    Sets attribute ``is_shape`` for outputs of a node.
+
+    :param node: node to process
+    :param values: stored results, values in this dictionary are updated
+    :param drop: variables not to consider because the come from the graph
+        holding this subgraph
+    :return: list of modified results
+    """
+    if not node.input:
+        # Constant
+        return []
+    drop = drop or set()
+    if node.op_type in ("Shape", "Size") and node.domain == "":
+        values[node.output[0]].is_shape = True
+        return [node.output[0]]
+    is_shapes = [values[i].is_shape for i in node.input if i not in drop]
+    if any(is_shapes):
+        if is_shapes[0] and len(node.output) == 1:
+            values[node.output[0]].is_shape = True
+            return [node.output[0]]
+    else:
+        for o in node.output:
+            values[o].is_shape = False
+        return list(node.output)
+    return []
+
+
+def first_used_last_used(
+    proto: Union[onnx.FunctionProto, onnx.GraphProto, onnx.ModelProto],
+    constant_as_initializer: bool = False,
+) -> Dict[str, RuntimeValue]:
+    """
+    Builds first used, last used information for every result
+    in the model.
+
+    :param proto: model, graph or function
+    :param constant_as_initializer: outputs of node Constant is tagged as INITIALIZER
+    :return: dictionary of RuntimeValue
+    """
+    values = {}
+    if isinstance(proto, onnx.ModelProto):
+        initializer = proto.graph.initializer
+        sparse_initializer = proto.graph.sparse_initializer
+        _input = proto.graph.input
+        output = proto.graph.output
+        _node = proto.graph.node
+        allow_unknown = False
+    elif isinstance(proto, onnx.GraphProto):
+        initializer = proto.initializer
+        sparse_initializer = proto.sparse_initializer
+        _input = proto.input
+        output = proto.output
+        _node = proto.node
+        allow_unknown = True
+    else:
+        initializer = []
+        sparse_initializer = []
+        _input = proto.input
+        output = proto.output
+        _node = proto.node
+        allow_unknown = False
+
+    for init in initializer:
+        values[init.name] = RuntimeValue(
+            init.name, kind=RuntimeValueKind.INITIALIZER, created=-1
+        )
+    for init in sparse_initializer:
+        values[init.name] = RuntimeValue(
+            init.name, created=-1, kind=RuntimeValueKind.INITIALIZER
+        )
+    for inp in _input:
+        n = inp if isinstance(inp, str) else inp.name
+        values[n] = RuntimeValue(n, created=-1, kind=RuntimeValueKind.INPUT)
+    drop = set()
+    for it, node in enumerate(_node):
+        for i in node.input:
+            if i not in values:
+                assert allow_unknown, f"Input {i!r} is unknown."
+                # This input comes from a context and the model is a GraphProto
+                drop.add(i)
+                continue
+            if values[i].first_used is None:
+                values[i].first_used = it
+            values[i].last_used = it
+        for att in node.attribute:
+            if att.type == onnx.AttributeProto.GRAPH:
+                for n in get_hidden_inputs(att.g):
+                    if values[n].first_used is None:
+                        values[n].first_used = it
+                    values[n].last_used = it
+        is_constant = node.op_type == "Constant" and node.domain == ""
+        for o in node.output:
+            values[o] = RuntimeValue(
+                o,
+                created=it,
+                kind=(
+                    RuntimeValueKind.INITIALIZER
+                    if is_constant and constant_as_initializer
+                    else RuntimeValueKind.RESULT
+                ),
+            )
+        set_is_shape(node, values, drop=drop)
+
+    for out in output:
+        n = out if isinstance(out, str) else out.name
+        values[n].kind = RuntimeValueKind.OUTPUT
+        values[n].last_used = len(_node)
+    return values

{onnx_diagnostic-0.6.0.dist-info → onnx_diagnostic-0.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-diagnostic
-Version: 0.6.0
+Version: 0.6.2
 Summary: Investigate ONNX models
 Home-page: https://github.com/sdpython/onnx-diagnostic
 Author: Xavier Dupré