onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/api.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+
+class TensorLike:
+    """Mocks a tensor."""
+
+    @property
+    def dtype(self) -> Any:
+        "Must be overwritten."
+        raise NotImplementedError("dtype must be overwritten.")
+
+    @property
+    def shape(self) -> Any:
+        "Must be overwritten."
+        raise NotImplementedError("shape must be overwritten.")

onnx_diagnostic/doc.py

@@ -0,0 +1,100 @@
+from typing import Optional
+import numpy as np
+
+
+def get_latest_pypi_version(package_name="onnx-diagnostic") -> str:
+    """Returns the latest published version."""
+
+    import requests
+
+    url = f"https://pypi.org/pypi/{package_name}/json"
+    response = requests.get(url)
+
+    assert response.status_code == 200, f"Unable to retrieve the version response={response}"
+    data = response.json()
+    version = data["info"]["version"]
+    return version
+
+
+def update_version_package(version: str, package_name="onnx-diagnostic") -> str:
+    "Adds dev if the major version is different from the latest published one."
+    released = get_latest_pypi_version(package_name)
+    shorten_r = ".".join(released.split(".")[:2])
+    shorten_v = ".".join(version.split(".")[:2])
+    return version if shorten_r == shorten_v else f"{shorten_v}.dev"
+
+
+def reset_torch_transformers(gallery_conf, fname):
+    "Resets torch dynamo for :epkg:`sphinx-gallery`."
+    import matplotlib.pyplot as plt
+    import torch
+
+    plt.style.use("ggplot")
+    torch._dynamo.reset()
+
+
+def plot_legend(
+    text: str, text_bottom: str = "", color: str = "green", fontsize: int = 15
+) -> "matplotlib.axes.Axes":  # noqa: F821
+    import matplotlib.pyplot as plt
+
+    fig = plt.figure(figsize=(2, 2))
+    ax = fig.add_subplot()
+    ax.axis([0, 5, 0, 5])
+    ax.text(2.5, 4, "END", fontsize=10, horizontalalignment="center")
+    ax.text(
+        2.5,
+        2.5,
+        text,
+        fontsize=fontsize,
+        bbox={"facecolor": color, "alpha": 0.5, "pad": 10},
+        horizontalalignment="center",
+        verticalalignment="center",
+    )
+    if text_bottom:
+        ax.text(4.5, 0.5, text_bottom, fontsize=7, horizontalalignment="right")
+    ax.grid(False)
+    ax.set_axis_off()
+    return ax
+
+
+def rotate_align(ax, angle=15, align="right"):
+    """Rotates x-label and align them to thr right. Returns ax."""
+    for label in ax.get_xticklabels():
+        label.set_rotation(angle)
+        label.set_horizontalalignment(align)
+    return ax
+
+
+def save_fig(ax, name: str):
+    """Applies ``tight_layout`` and saves the figures. Returns ax."""
+    import matplotlib.pyplot as plt
+
+    plt.tight_layout()
+    fig = ax.get_figure()
+    fig.savefig(name)
+    return ax
+
+
+def title(ax: "plt.axes", title: str) -> "plt.axes":  # noqa: F821
+    "Adds a title to axes and returns them."
+    ax.set_title(title)
+    return ax
+
+
+def plot_histogram(
+    tensor: np.ndarray,
+    ax: Optional["plt.axes"] = None,  # noqa: F821
+    bins: int = 30,
+    color: str = "orange",
+    alpha: float = 0.7,
+) -> "plt.axes":  # noqa: F821
+    "Computes the distribution for a tensor."
+    if ax is None:
+        import matplotlib.pyplot as plt
+
+        ax = plt.gca()
+        ax.cla()
+    ax.hist(tensor, bins=30, color="orange", alpha=0.7)
+    ax.set_yscale("log")
+    return ax

onnx_diagnostic/export/__init__.py

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

onnx_diagnostic/export/api.py

@@ -0,0 +1,124 @@
+from typing import Any, Dict, List, Sequence, Optional, Tuple, Union
+import torch
+
+
+def to_onnx(
+    mod: Union["torch.nn.Module", "torch.fx.GraphModule"],  # noqa: F821
+    args: Optional[Sequence["torch.Tensor"]] = None,  # noqa: F821
+    kwargs: Optional[Dict[str, "torch.Tensor"]] = None,  # noqa: F821
+    input_names: Optional[Sequence[str]] = None,
+    target_opset: Optional[Union[int, Dict[str, int]]] = None,
+    verbose: int = 0,
+    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+    filename: Optional[str] = None,
+    output_names: Optional[List[str]] = None,
+    output_dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+    exporter: str = "onnx-dynamo",
+) -> Any:
+    """
+    Common API for exporters. By default, the models are optimized to use the
+    most efficient kernels implemented in :epkg:`onnxruntime`.
+
+    :param mod: torch model
+    :param args: unnamed arguments
+    :param kwargs: named arguments
+    :param input_names: input names for the onnx model (optional)
+    :param target_opset: opset to target, if not specified, each converter
+        keeps its default value
+    :param verbose: verbosity level
+    :param dynamic_shapes: dynamic shapes, usually a nested structure
+        included a dictionary for each tensor
+    :param filename: output filename
+    :param output_names: to change the output of the onnx model
+    :param output_dynamic_shapes: to overwrite the dynamic shapes names
+    :param exporter: exporter to use (``onnx-dynamo``, ``modelbuilder``, ``custom``)
+    :return: the output of the selected exporter, usually a structure including
+        an onnx model
+
+    A simple example:
+
+    .. code-block:: python
+
+        to_onnx(
+            model,
+            kwargs=inputs,
+            dynamic_shapes=ds,
+            exporter=exporter,
+            filename=filename,
+        )
+    """
+    if exporter == "custom":
+        from experimental_experiment.torch_interpreter import to_onnx as _to_onnx
+        from experimental_experiment.xbuilder import OptimizationOptions
+
+        return _to_onnx(
+            mod,
+            args=args,
+            kwargs=kwargs,
+            input_names=input_names,
+            output_names=output_names,
+            target_opset=target_opset,
+            verbose=verbose,
+            filename=filename,
+            dynamic_shapes=dynamic_shapes,
+            large_model=True,
+            output_dynamic_shapes=output_dynamic_shapes,
+            options=OptimizationOptions(patterns="default+onnxruntime"),
+        )
+    if exporter in ("dynamo", "onnx-dynamo"):
+        import onnxscript.rewriter.ort_fusions as ort_fusions
+
+        assert (
+            not output_dynamic_shapes
+        ), f"output_dynamic_shapes not supported for exporter={exporter!r}"
+        epo = torch.onnx.export(
+            mod,
+            args=args or tuple(),
+            kwargs=kwargs,
+            input_names=input_names,
+            output_names=output_names,
+            opset_version=target_opset,
+            dynamic_shapes=dynamic_shapes,
+            dynamo=True,
+        )
+        ort_fusions.optimize_for_ort(epo.model)
+        epo.save(filename)
+        return epo
+
+    if exporter == "modelbuilder":
+        import os
+        from ..helpers import flatten_object, string_type
+        from ..helpers.model_builder_helper import create_model_builder, save_model_builder
+
+        assert filename, f"filename must be specified for exporter={exporter!r}"
+        assert (
+            not output_dynamic_shapes
+        ), f"output_dynamic_shapes not supported for exporter={exporter!r}"
+        assert hasattr(mod, "config"), f"configuration is missing in model class {type(mod)}"
+        assert not args, f"only kwargs can be defined with exporter={exporter!r}"
+        assert list(kwargs) == ["input_ids", "attention_mask", "past_key_values"], (  # type: ignore[arg-type]
+            f"Only a specified set of inputs is supported for exporter={exporter!r}, "
+            f"but it is {list(kwargs)}"  # type: ignore[arg-type]
+        )
+        flat_inputs = flatten_object(kwargs, drop_keys=True)
+        first = flat_inputs[0]
+        first_float = [
+            t
+            for t in flat_inputs
+            if t.dtype in {torch.float32, torch.double, torch.float16, torch.bfloat16}
+        ]
+        assert first_float, (
+            f"Unable to find a float tensor in the inputs "
+            f"{string_type(kwargs, with_shape=True)}"
+        )
+        onx = create_model_builder(
+            mod.config,
+            mod,
+            precision=str(first_float[0].dtype).split(".")[-1],
+            execution_provider="cuda" if first.is_cuda else "cpu",
+            cache_dir=os.path.dirname(filename),
+        )
+        save_model_builder(onx, os.path.dirname(filename))
+        return onx
+
+    raise ValueError(f"Unknown exporter={exporter!r}")