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. Patches can be enabled as follows:
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.
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
Torch Export
Use DYNAMIC or AUTO when exporting if dynamic shapes has constraints
Steel method forward to guess the dynamic shapes (with Tiny-LLM)
Investigate ONNX models
Snapshot of usefuls tools
torch_export_patches
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
from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
dynamic_shapes = all_dynamic_shapes_from_inputs(cache)torch_export_rewrite
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
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
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
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
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'}), {})
1 maintainer
Metadata
1 maintainer
Release files for onnx-diagnostic 0.9.3
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| onnx_diagnostic-0.9.3-py3-none-any.whl | Python 3 | none | any | Details |
Release files / onnx_diagnostic-0.9.3-py3-none-any.whl
| Download URL | onnx_diagnostic-0.9.3-py3-none-any.whl |
|---|---|
| Size | 1.9 MB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
ff446b7670138346113ef766089cff2d9c4683ad48b95a789af34929a24cf98f
|
|
BLAKE2b-256 checksum How to use checksums |
3b618286b57c50c66635bdb6a7362d770aa8ae2d626e961558926986c2159612
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.12.3
|
