onnx-diagnostic: investigate onnx models¶

The main feature is about patches: it helps exporting pytorch models into ONNX, mostly designed for LLMs using dynamic caches. Sources available at github/onnx-diagnostic.

with torch_export_patches(patch_transformers=True) as f:
    ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
    # ...

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

Contents

git clone https://github.com/sdpython/onnx-diagnostic.git
cd onnx-diagnostic
pip install -e .

pip install onnx-diagnostic

Enlightening Examples¶

Where to start to export a model

Export microsoft/phi-2

Torch Export

Investigate ONNX models

Some Usefuls Tools¶

torch_export_patches¶

See onnx_diagnostic.torch_export_patches.torch_export_patches().

with torch_export_patches(patch_transformers=True) as f:
    ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
    # ...

torch_export_rewrite¶

See onnx_diagnostic.torch_export_patches.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¶

See onnx_diagnostic.helpers.string_type().

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¶

See onnx_diagnostic.helpers.onnx_helper.onnx_dtype_name().

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¶

See onnx_diagnostic.helpers.max_diff().

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¶

See onnx_diagnostic.export.ModelInputs.guess_dynamic_shapes().

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'}), {})

use_dyn_not_str¶

See onnx_diagnostic.torch_export_patches.patch_inputs.use_dyn_not_str().

The function replaces dynamic dimensions defined as strings by torch.export.Dim.DYNAMIC.

Older versions¶

The documentation was updated on:

    2025-06-06 15:21:03.754762

With the following versions:

    numpy: 2.2.6
    ml_dtypes: 0.5.1
    sklearn: 1.6.1
    onnx: 1.19.0
    onnxruntime: 1.23.0+cu126
    onnxscript: 0.3.0.dev20250301
    torch: 2.8.0.dev20250603+cu126
    transformers: 4.53.0.dev0
    timm: 1.0.15
    has_onnxruntime_training: True

Size of the package:

                              lines   chars
    ext dir                                
    .py                        1564   43553
        export                  963   30114
        helpers                5416  161831
        reference              3059   91721
        tasks                  1450   46792
        torch_export_patches   3975  128471
        torch_models           5994  290675
        torch_onnx              591   17414