https://github.com/sdpython/onnx-diagnostic

Investigate onnx models
https://github.com/sdpython/onnx-diagnostic

Last synced: about 2 months ago
JSON representation

Investigate onnx models

• Host: GitHub
• URL: https://github.com/sdpython/onnx-diagnostic
• Owner: sdpython
• License: mit
• Created: 2025-03-21T08:59:40.000Z (12 months ago)
• Default Branch: main
• Last Pushed: 2026-01-18T12:40:03.000Z (about 2 months ago)
• Last Synced: 2026-01-18T14:13:54.752Z (about 2 months ago)
• Language: Python
• Size: 2.37 MB
• Stars: 4
• Watchers: 1
• Forks: 1
• Open Issues: 6
• Metadata Files:
  • Readme: README.rst
  • Changelog: CHANGELOGS.rst
  • License: LICENSE.txt

Awesome Lists containing this project

README

          

.. image:: https://github.com/sdpython/onnx-diagnostic/raw/main/_doc/_static/logo.png

    :width: 120

onnx-diagnostic: investigate onnx models

========================================

.. image:: https://github.com/sdpython/onnx-diagnostic/actions/workflows/documentation.yml/badge.svg

    :target: https://github.com/sdpython/onnx-diagnostic/actions/workflows/documentation.yml

.. image:: https://badge.fury.io/py/onnx-diagnostic.svg

    :target: http://badge.fury.io/py/onnx-diagnostic

.. image:: https://img.shields.io/badge/license-MIT-blue.svg

    :alt: MIT License

    :target: https://opensource.org/license/MIT/

.. image:: https://img.shields.io/github/repo-size/sdpython/onnx-diagnostic

    :target: https://github.com/sdpython/onnx-diagnostic/

    :alt: size

.. image:: https://img.shields.io/badge/code%20style-black-000000.svg

    :target: https://github.com/psf/black

.. image:: https://codecov.io/gh/sdpython/onnx-diagnostic/graph/badge.svg?token=91T5ZVIP96 

    :target: https://codecov.io/gh/sdpython/onnx-diagnostic

The main feature is about `patches `_:

it helps exporting **pytorch models into ONNX**, mostly designed for LLMs using dynamic caches.

Patches can be enabled as follows:

.. code-block:: python

  from onnx_diagnostic.torch_export_patches import torch_export_patches

  with torch_export_patches(patch_transformers=True) as f:

      ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)

      # ...

Dynamic shapes are difficult to guess for caches, one function

returns a structure defining all dimensions as dynamic.

You need then to remove those which are not dynamic in your model.

.. code-block:: python

  from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs

  dynamic_shapes = all_dynamic_shapes_from_inputs(cache)

It also implements tools to investigate, validate exported models (ExportedProgramm, ONNXProgram, ...).

See `documentation of onnx-diagnostic `_ and

`torch_export_patches `_.

Getting started

+++++++++++++++

::

    git clone https://github.com/sdpython/onnx-diagnostic.git

    cd onnx-diagnostic

    pip install -e . -v

or

::

    pip install onnx-diagnostic

Enlightening Examples

+++++++++++++++++++++

**Where to start to export a model**

* `Export microsoft/phi-2

  `_

* `Export a LLM through method generate (with Tiny-LLM)

  `_

**Torch Export**

* `Use DYNAMIC or AUTO when exporting if dynamic shapes has constraints

  `_

* `Find and fix an export issue due to dynamic shapes

  `_

* `Export with DynamicCache and guessed dynamic shapes

  `_

* `Steel method forward to guess the dynamic shapes (with Tiny-LLM)

  `_

* `Export Tiny-LLM with patches

  `_

**Investigate ONNX models**

* `Find where a model is failing by running submodels

  `_

* `Intermediate results with (ONNX) ReferenceEvaluator

  `_

* `Intermediate results with onnxruntime

  `_

Snapshot of usefuls tools

+++++++++++++++++++++++++

**torch_export_patches**

.. code-block:: python

  from onnx_diagnostic.torch_export_patches import torch_export_patches

  with torch_export_patches(patch_transformers=True) as f:

      ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)

      # ...

**all_dynamic_shapes_from_inputs**

.. code-block:: python

  from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs

  dynamic_shapes = all_dynamic_shapes_from_inputs(cache)

**torch_export_rewrite**

.. code-block:: python

  from onnx_diagnostic.torch_export_patches import torch_export_rewrite

  with torch_export_rewrite(rewrite=[Model.forward]) as f:

      ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)

      # ...

**string_type**

.. code-block:: python

    import torch

    from onnx_diagnostic.helpers import string_type

    inputs = (

        torch.rand((3, 4), dtype=torch.float16),

        [torch.rand((5, 6), dtype=torch.float16), torch.rand((5, 6, 7), dtype=torch.float16)],

    )

    # with shapes

    print(string_type(inputs, with_shape=True))

::

    >>> (T10s3x4,#2[T10s5x6,T10s5x6x7])

**onnx_dtype_name**

.. code-block:: python

        import onnx

        from onnx_diagnostic.helpers.onnx_helper import onnx_dtype_name

        itype = onnx.TensorProto.BFLOAT16

        print(onnx_dtype_name(itype))

        print(onnx_dtype_name(7))

::

    >>> BFLOAT16

    >>> INT64

**max_diff**

.. code-block:: python

    import torch

    from onnx_diagnostic.helpers import max_diff

    print(

        max_diff(

            (torch.Tensor([1, 2]), (torch.Tensor([1, 2]),)),

            (torch.Tensor([1, 2]), (torch.Tensor([1, 2]),)),

        )

    )

::

    >>> {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 4.0, "dnan": 0.0}s

**guess_dynamic_shapes**

.. code-block:: python

    inputs = [

        (torch.randn((5, 6)), torch.randn((1, 6))),

        (torch.randn((7, 8)), torch.randn((1, 8))),

    ]

    ds = ModelInputs(model, inputs).guess_dynamic_shapes(auto="dim")

    print(ds)

::

    >>> (({0: 'dim_0I0', 1: 'dim_0I1'}, {1: 'dim_1I1'}), {})