Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,26 @@ onnx-diagnostic: investigate onnx models

The main feature is about `patches <https://github.com/sdpython/onnx-diagnostic/tree/main/onnx_diagnostic/torch_export_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_shape_from_inputs

dynamic_shapes = all_dynamic_shape_from_inputs(cache)

It also implements tools to investigate, validate exported models (ExportedProgramm, ONNXProgram, ...).
See `documentation of onnx-diagnostic <https://sdpython.github.io/doc/onnx-diagnostic/dev/>`_ and
`torch_export_patches <https://sdpython.github.io/doc/onnx-diagnostic/dev/api/torch_export_patches/index.html#onnx_diagnostic.torch_export_patches.torch_export_patches>`_.
Expand Down Expand Up @@ -90,14 +103,26 @@ Snapshot of usefuls tools

.. 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_shape_from_inputs**

.. code-block:: python

from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs

dynamic_shapes = all_dynamic_shape_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)
# ...
Expand Down
24 changes: 24 additions & 0 deletions _doc/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,26 @@ onnx-diagnostic: investigate onnx models
The main feature is about `patches <https://github.com/sdpython/onnx-diagnostic/tree/main/onnx_diagnostic/torch_export_patches>`_:
it helps exporting **pytorch models into ONNX**, mostly designed for LLMs using dynamic caches.
Sources available at `github/onnx-diagnostic <https://github.com/sdpython/onnx-diagnostic/>`_.
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_shape_from_inputs

dynamic_shapes = all_dynamic_shape_from_inputs(cache)

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

Expand Down Expand Up @@ -119,6 +132,17 @@ See :func:`onnx_diagnostic.torch_export_patches.torch_export_rewrite`.
ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
# ...

all_dynamic_shape_from_inputs
+++++++++++++++++++++++++++++

See :func:`onnx_diagnostic.export.shape_helper.all_dynamic_shape_from_inputs`.

.. code-block:: python

from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs

dynamic_shapes = all_dynamic_shape_from_inputs(cache)

string_type
+++++++++++

Expand Down
Loading