doc (#78)

* doc

* improve documentation

* fix doc

* fix a few things

* update CI

* onnx

Author: xadupre

.github/workflows/ci.yml

@@ -64,6 +64,11 @@ jobs:
       - name: Install requirements dev
         run: python -m pip install -r requirements-dev.txt
 
+      - name: Uninstall onnx and install onnx-weekly
+        run: |
+          python -m pip uninstall -y onnx
+          python -m pip install onnx-weekly
+
       - name: Cache pip
         uses: actions/cache@v4
         with:

.github/workflows/documentation.yml

@@ -41,6 +41,11 @@ jobs:
       - name: Install requirements dev
         run: python -m pip install -r requirements-dev.txt
 
+      - name: Uninstall onnx and install onnx-weekly
+        run: |
+          python -m pip uninstall -y onnx
+          python -m pip install onnx-weekly
+
       - name: Cache pip
         uses: actions/cache@v4
         with:

.gitignore

@@ -37,6 +37,7 @@ dump_models/*
 dump_bash_bench/*
 dump_llama/*
 dump_test*
+dump_validate*
 dump_sdpa_*
 temp_dump_models/*
 dump_dort_bench/*

README.rst

@@ -56,6 +56,11 @@ or
 Enlightening Examples
 +++++++++++++++++++++
 
+**Where to start to export a model**
+
+* `Export microsoft/phi-2
+  <https://sdpython.github.io/doc/onnx-diagnostic/dev/auto_examples/plot_export_tiny_phi2.html>`_
+
 **Torch Export**
 
 * `Use DYNAMIC or AUTO when exporting if dynamic shapes has constraints

_doc/examples/plot_export_tiny_llm.py

@@ -1,11 +1,11 @@
 """
 .. _l-plot-tiny-llm-export:
 
-Steel method forward to guess the dynamic shapes (with Tiny-LLM)
-================================================================
+Steel method forward to guess inputs and dynamic shapes (with Tiny-LLM)
+=======================================================================
 
 Inputs are always dynamic with LLMs that is why dynamic shapes
-needs to be specified when a LLM is exported with:func:`torch.export.export`.
+needs to be specified when a LLM is exported with :func:`torch.export.export`.
 Most of the examples on :epkg:`HuggingFace` use method
 :meth:`transformers.GenerationMixin.generate` but we only want to
 export the model and its method ``forward``.

_doc/examples/plot_export_tiny_phi2.py

@@ -1,11 +1,16 @@
 """
 .. _l-plot-export_tiny_phi2:
 
-Untrained microsoft/phi-2
-=========================
+======================
+Export microsoft/phi-2
+======================
 
-:epkg:`microsoft/phi-2` is not a big models but still quite big
-when it comes to write unittest. Function
+This function exports an smaller untrained model with the same architecture.
+It is faster than the pretrained model.
+When this works, the untrained model can be replaced by the trained one.
+
+:epkg:`microsoft/phi-2` is not a big model but still quite big
+when it comes to write unittests. Function
 :func:`onnx_diagnostic.torch_models.hghub.get_untrained_model_with_inputs`
 can be used to create a reduced untrained version of a model coming from
 :epkg:`HuggingFace`. It downloads the configuration from the website
@@ -14,7 +19,7 @@
 the export or to compare performance. The relevance does not matter.
 
 Create the dummy model
-++++++++++++++++++++++
+======================
 """
 
 import copy
@@ -45,9 +50,11 @@
     data["n_weights"],
 )
 
-print(f"model {size / 2**20:1.3f} Mb with {n_weights // 1000} mille parameters.")
+print(f"model {size / 2**20:1.1f} Mb with {n_weights // 1000} thousands of parameters.")
 # %%
 # The original model has 2.7 billion parameters. It was divided by more than 10.
+# However, it can still be used with
+# ``get_untrained_model_with_inputs("microsoft/phi-2", same_as_pretrained=True)``.
 # Let's see the configuration.
 print(config)
 
@@ -72,13 +79,18 @@
 
 
 # %%
-# Export
-# ++++++
+# Export to fx.Graph
+# ==================
+#
+# :func:`torch.export.export` is the first step before converting
+# a model into ONNX. The inputs are duplicated (with ``copy.deepcopy``)
+# because the model may modify them inline (a cache for example).
+# Shapes may not match on the second call with the modified inputs.
 
 
-with torch_export_patches(patch_transformers=True) as modificator:
+with torch_export_patches(patch_transformers=True):
 
-    # Unnecessary steps but useful in case of an error
+    # Two unnecessary steps but useful in case of an error
     # We check the cache is registered.
     assert is_cache_dynamic_registered()
 
@@ -88,24 +100,26 @@
         d["abs"] < 1e-5
     ), f"The model with patches produces different outputs: {string_diff(d)}"
 
-    # Then we export.
+    # Then we export: the only import line in this section.
     ep = torch.export.export(
         untrained_model,
         (),
-        kwargs=modificator(copy.deepcopy(inputs)),
+        kwargs=copy.deepcopy(inputs),
         dynamic_shapes=use_dyn_not_str(dynamic_shapes),
         strict=False,  # mandatory for torch==2.6
     )
 
     # We check the exported program produces the same results as well.
+    # This step is again unnecessary.
     d = max_diff(expected, ep.module()(**copy.deepcopy(inputs)))
     assert d["abs"] < 1e-5, f"The exported model different outputs: {string_diff(d)}"
 
 # %%
 # Export to ONNX
-# ++++++++++++++
+# ==============
 #
-# The export works. We can export to ONNX now.
+# The export works. We can export to ONNX now
+# :func:`torch.onnx.export`.
 # Patches are still needed because the export
 # applies :meth:`torch.export.ExportedProgram.run_decompositions`
 # may export local pieces of the model again.
@@ -157,4 +171,52 @@
 # It looks good.
 
 # %%
-doc.plot_legend("untrained smaller\nmicrosoft/phi-2", "torch.onnx.export", "orange")
+doc.plot_legend("export\nuntrained smaller\nmicrosoft/phi-2", "torch.onnx.export", "orange")
+
+# %%
+# Possible Issues
+# ===============
+#
+# Unknown task
+# ++++++++++++
+#
+# Function :func:`onnx_diagnostic.torch_models.hghub.get_untrained_model_with_inputs`
+# is unabl to guess a task associated to the model.
+# A different set of dummy inputs is defined for every task.
+# The user needs to explicitly give that information to the function.
+# Tasks are the same as the one defined by
+# `HuggingFace/models <https://huggingface.co/models>`_.
+#
+# Inputs are incorrect
+# ++++++++++++++++++++
+#
+# Example :ref:`l-plot-tiny-llm-export` explains
+# how to retrieve that information. If you cannot guess the dynamic
+# shapes - a cache can be tricky sometimes, follow example
+# :ref:`l-plot-export-with-args-kwargs`.
+#
+# DynamicCache or any other cache cannot be exported
+# ++++++++++++++++++++++++++++++++++++++++++++++++++
+#
+# That's the role of :func:`onnx_diagnostic.torch_export_patches.torch_export_patches`.
+# It registers the necessary information into pytorch to make the export
+# work with these. Its need should slowly disappear until :epkg:`transformers`
+# includes the serialization functions.
+#
+# Control Flow
+# ++++++++++++
+#
+# Every mixture of models goes through a control flow (a test).
+# It also happens when a cache is truncated. The code of the model
+# needs to be changed. See example :ref:`l-plot-export-cond`.
+# Loops are not supported yet.
+#
+# Issue with dynamic shapes
+# +++++++++++++++++++++++++
+#
+# Example :ref:`l-plot-dynamic-shapes-python-int` gives one reason
+# this process may fail but that's not the only one.
+# Example :ref:`l-plot-export-locale-issue` gives an way to locate
+# the cause but that does not cover all the possible causes.
+# Raising an issue on github would be the recommended option
+# until it is fixed.

_doc/index.rst

@@ -65,6 +65,10 @@ or
 Enlightening Examples
 +++++++++++++++++++++
 
+**Where to start to export a model**
+
+* :ref:`l-plot-export_tiny_phi2`
+
 **Torch Export**
 
 * :ref:`l-plot-export-cond`

_doc/recipes/plot_dynamic_shapes_python_int.py

@@ -1,4 +1,6 @@
 """
+.. _l-plot-dynamic-shapes-python-int:
+
 Do not use python int with dynamic shapes
 =========================================
 

_unittests/ut_reference/test_reference_back.py

@@ -15,6 +15,7 @@
 
 class ExtendedReferenceEvaluatorBackendRep(onnx.backend.base.BackendRep):
     def __init__(self, session):
+        super().__init__()
         self._session = session
 
     def run(self, inputs, **kwargs):

onnx_diagnostic/tasks/image_classification.py

@@ -7,6 +7,13 @@
 
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
+    if (
+        hasattr(config, "model_type")
+        and config.model_type == "timm_wrapper"
+        and not hasattr(config, "num_hidden_layers")
+    ):
+        # We cannot reduce.
+        return {}
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
@@ -82,6 +89,20 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     If the configuration is None, the function selects typical dimensions.
     """
     if config is not None:
+        if (
+            hasattr(config, "model_type")
+            and config.model_type == "timm_wrapper"
+            and not hasattr(config, "num_hidden_layers")
+        ):
+            input_size = config.pretrained_cfg["input_size"]
+            kwargs = dict(
+                batch_size=2,
+                input_width=input_size[-2],
+                input_height=input_size[-1],
+                input_channels=input_size[-3],
+            )
+            return kwargs, get_inputs
+
         check_hasattr(config, ("image_size", "architectures"), "num_channels")
     if config is not None:
         if hasattr(config, "image_size"):