diff --git a/README.rst b/README.rst index 34b7f2ee..c876fbb2 100644 --- a/README.rst +++ b/README.rst @@ -27,13 +27,26 @@ 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: .. 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 `_ and `torch_export_patches `_. @@ -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) # ... diff --git a/_doc/index.rst b/_doc/index.rst index 322c25cb..8cdfd4e6 100644 --- a/_doc/index.rst +++ b/_doc/index.rst @@ -21,13 +21,26 @@ 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 `_. +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`. @@ -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 +++++++++++