uses onnx-weekly on CI, fix documentation

.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"):

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -88,6 +88,15 @@ def get_model_info(model_id) -> Any:
     return model_info(model_id)
 
 
+def _guess_task_from_config(config: Any) -> Optional[str]:
+    """Tries to infer a task from the configuration."""
+    if hasattr(config, "bbox_loss_coefficient") and hasattr(config, "giou_loss_coefficient"):
+        return "object-detection"
+    if hasattr(config, "architecture") and config.architecture:
+        return task_from_arch(config.architecture)
+    return None
+
+
 @functools.cache
 def task_from_arch(arch: str, default_value: Optional[str] = None) -> str:
     """
@@ -126,7 +135,7 @@ def task_from_id(
     :param default_value: if specified, the function returns this value
         if the task cannot be determined
     :param pretrained: uses the config
-    :param fall_back_to_pretrained: balls back to pretrained config
+    :param fall_back_to_pretrained: falls back to pretrained config
     :return: task
     """
     if not pretrained:
@@ -139,6 +148,9 @@ def task_from_id(
     try:
         return config.pipeline_tag
     except AttributeError:
+        guess = _guess_task_from_config(config)
+        if guess is not None:
+            return guess
         assert config.architectures is not None and len(config.architectures) == 1, (
             f"Cannot return the task of {model_id!r}, pipeline_tag is not setup, "
             f"architectures={config.architectures} in config={config}"

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -78,6 +78,7 @@
     MobileBertModel,feature-extraction
     MobileNetV1Model,image-feature-extraction
     MobileNetV2Model,image-feature-extraction
+    mobilenetv3_small_100,image-classification
     MobileViTForImageClassification,image-classification
     ModernBertForMaskedLM,fill-mask
     Phi4MMForCausalLM,MoE
@@ -202,7 +203,7 @@ def load_models_testing() -> List[str]:
 @functools.cache
 def load_architecture_task() -> Dict[str, str]:
     """
-    Returns a dictionary mapping architecture to task.
+    Returns a dictionary mapping architectures to tasks.
 
     import pprint
     from onnx_diagnostic.torch_models.hghub.hub_data import load_architecture_task

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -4,7 +4,7 @@
 import transformers
 from ...helpers.config_helper import update_config
 from ...tasks import reduce_model_config, random_input_kwargs
-from .hub_api import task_from_arch, get_pretrained_config
+from .hub_api import task_from_arch, task_from_id, get_pretrained_config
 
 
 def get_untrained_model_with_inputs(
@@ -64,17 +64,21 @@ def get_untrained_model_with_inputs(
         config = get_pretrained_config(
             model_id, use_preinstalled=use_preinstalled, **(model_kwargs or {})
         )
+    if hasattr(config, "architecture") and config.architecture:
+        archs = [config.architecture]
     archs = config.architectures  # type: ignore
-    assert archs is not None and len(archs) == 1, (
+    task = None
+    if archs is None:
+        task = task_from_id(model_id)
+    assert task is not None or (archs is not None and len(archs) == 1), (
         f"Unable to determine the architecture for model {model_id!r}, "
         f"architectures={archs!r}, conf={config}"
     )
-    arch = archs[0]
-    if verbose:
-        print(f"[get_untrained_model_with_inputs] architecture={arch!r}")
     if verbose:
+        print(f"[get_untrained_model_with_inputs] architectures={archs!r}")
         print(f"[get_untrained_model_with_inputs] cls={config.__class__.__name__!r}")
-    task = task_from_arch(arch)
+    if task is None:
+        task = task_from_arch(archs[0])
     if verbose:
         print(f"[get_untrained_model_with_inputs] task={task!r}")
 
@@ -106,7 +110,15 @@ def get_untrained_model_with_inputs(
     if inputs_kwargs:
         kwargs.update(inputs_kwargs)
 
-    model = getattr(transformers, arch)(config)
+    if archs is not None:
+        model = getattr(transformers, archs[0])(config)
+    else:
+        assert same_as_pretrained, (
+            f"Model {model_id!r} cannot be built, the model cannot be built. "
+            f"It must be downloaded. Use same_as_pretrained=True."
+        )
+        model = None
+
     # This line is important. Some models may produce different
     # outputs even with the same inputs in training mode.
     model.eval()

requirements-dev.txt

@@ -5,6 +5,7 @@ furo
 huggingface_hub
 matplotlib
 onnx-array-api>=0.3.1
+onnx
 git+https://github.com/microsoft/onnxscript.git
 openpyxl
 packaging