add option drop_names

Author: xadupre

.github/workflows/check-urls.yml

@@ -43,5 +43,5 @@ jobs:
         timeout: 2
         retry_count# : 2
         exclude_urls: https://hal.archives-,ouvertes.fr/hal-00990252/document,http://badge.fury.io/py/onnx-diagnostic,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://medium.com/@msouza.os/llm-from-scratch-with-pytorch-9f21808c6319,https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/symbolic_shapes.py#L5965,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/huggingface/transformers/pull/36311
-        exclude_patterns: https://www.data.gouv.fr/fr/datasets/r/e3d83ab3-dc52-4c99-abaf-8a38050cc68c,https://dev.azure.com/,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://github.com/pytorch/pytorch/blob/main/torch/,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/huggingface/transformers/pull/36311
+        exclude_patterns: https://www.data.gouv.fr/fr/datasets/r/e3d83ab3-dc52-4c99-abaf-8a38050cc68c,https://dev.azure.com/,https://azure.microsoft.com/en-us/products/devops/pipelines,https://github.com/pytorch/pytorch/pull/83137#issuecomment-1211320670,https://github.com/NVIDIA/TransformerEngine.git@6a9edc38bf9b941b7d369af5103fa8fe0b121d61,https://github.com/pytorch/pytorch/blob/main/torch/,https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-04.html,https://badge.fury.io/py/onnx-diagnostic.svg,https://github.com/huggingface/transformers/pull/36311,https://codecov.io/
         # force_pass : true

_unittests/ut_torch_models/test_test_helpers.py

@@ -1,7 +1,11 @@
 import copy
 import unittest
 from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, ignore_warnings
-from onnx_diagnostic.torch_models.test_helper import get_inputs_for_task, validate_model
+from onnx_diagnostic.torch_models.test_helper import (
+    get_inputs_for_task,
+    validate_model,
+    filter_inputs,
+)
 from onnx_diagnostic.torch_models.hghub.model_inputs import get_get_inputs_function_for_tasks
 
 
@@ -63,6 +67,27 @@ def test_validate_model_onnx(self):
         self.assertIsInstance(data, dict)
         self.assertLess(summary["disc_onnx_ort_run_abs"], 1e-4)
 
+    def test_filter_inputs(self):
+        inputs, ds = {"a": 1, "b": 2}, {"a": 20, "b": 30}
+        ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["a"])
+        self.assertEqual((ni, nd), ({"b": 2}, {"b": 30}))
+
+        inputs, ds = (1, 2), {"a": 20, "b": 30}
+        ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["b"], model=["a", "b"])
+        self.assertEqual((ni, nd), ((1, None), {"a": 20}))
+
+        inputs, ds = (1, 2), (20, 30)
+        ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["b"], model=["a", "b"])
+        self.assertEqual((ni, nd), ((1, None), (20, None)))
+
+        inputs, ds = ((1,), {"b": 4}), {"a": 20, "b": 30}
+        ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["b"], model=["a", "b"])
+        self.assertEqual((ni, nd), ((1,), {"a": 20}))
+
+        inputs, ds = ((1,), {"b": 4}), {"a": 20, "b": 30}
+        ni, nd = filter_inputs(inputs, dynamic_shapes=ds, drop_names=["a"], model=["a", "b"])
+        self.assertEqual((ni, nd), (((None,), {"b": 4}), {"b": 30}))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/_command_lines_parser.py

@@ -276,6 +276,11 @@ def get_parser_validate() -> ArgumentParser:
         help="if not empty, a folder is created to dumps statistics, "
         "exported program, onnx...",
     )
+    parser.add_argument(
+        "--drop",
+        help="drops the following inputs names, it should be a list "
+        "with comma separated values",
+    )
     parser.add_argument("-v", "--verbose", default=0, type=int, help="verbosity")
     parser.add_argument("--dtype", help="changes dtype if necessary")
     parser.add_argument("--device", help="changes the device if necessary")
@@ -317,6 +322,7 @@ def _cmd_validate(argv: List[Any]):
             optimization=args.opt,
             exporter=args.export,
             dump_folder=args.dump_folder,
+            drop_inputs=None if not args.drop else args.drop.split(","),
         )
         print("")
         print("-- summary --")

onnx_diagnostic/torch_models/test_helper.py

@@ -1,5 +1,6 @@
+import inspect
 import os
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 import time
 import torch
 from ..helpers import max_diff, string_type, string_diff
@@ -45,6 +46,69 @@ def get_inputs_for_task(task: str, config: Optional[Any] = None) -> Dict[str, An
     return f(model=None, config=config, **kwargs)
 
 
+def split_args_kwargs(inputs: Any) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
+    """Splits into args, kwargs."""
+    if isinstance(inputs, dict):
+        return (), inputs
+    if isinstance(inputs, tuple) and len(inputs) == 2 and isinstance(inputs[1], dict):
+        return inputs
+    assert isinstance(inputs, tuple), f"Unexpected inputs {string_type(inputs)}"
+    return inputs, {}
+
+
+def make_inputs(
+    args: Optional[Tuple[Any, ...]], kwargs: Optional[Dict[str, Any]] = None
+) -> Any:
+    """Returns either args, kwargs or both depending on which ones are empty."""
+    assert args or kwargs, "No input was given."
+    if not args:
+        return kwargs
+    if not kwargs:
+        return args
+    return args, kwargs
+
+
+def filter_inputs(
+    inputs: Any,
+    drop_names: List[str],
+    model: Optional[Union[torch.nn.Module, List[str]]] = None,
+    dynamic_shapes: Optional[Any] = None,
+):
+    """
+    Drops some inputs from the given inputs.
+    It updates the dynamic shapes as well.
+    """
+    args, kwargs = split_args_kwargs(inputs)
+    set_drop_names = set(drop_names)
+    kwargs = {k: v for k, v in kwargs.items() if k not in set_drop_names}
+    dyn = (
+        {k: v for k, v in dynamic_shapes.items() if k not in set_drop_names}
+        if dynamic_shapes and isinstance(dynamic_shapes, dict)
+        else dynamic_shapes
+    )
+    if not args or all(i in kwargs for i in set_drop_names):
+        return make_inputs(args, kwargs), dyn
+    assert model, (
+        f"we need the model to get the parameter name but model is None, "
+        f"input_names={drop_names} and args={string_type(args)}"
+    )
+    pnames = (
+        list(inspect.signature(model.forward).parameters)
+        if isinstance(model, torch.nn.Module)
+        else model
+    )
+    new_args = []
+    new_ds = []
+    for i, a in enumerate(args):
+        if isinstance(dynamic_shapes, tuple):
+            new_ds.append(None if pnames[i] in set_drop_names else dynamic_shapes[i])
+        new_args.append(None if pnames[i] in set_drop_names else a)
+    new_inputs = make_inputs(tuple(new_args), kwargs)
+    if new_ds:
+        return new_inputs, tuple(new_ds)
+    return new_inputs, dyn
+
+
 def validate_model(
     model_id: str,
     task: Optional[str] = None,
@@ -59,6 +123,7 @@ def validate_model(
     quiet: bool = False,
     patch: bool = False,
     dump_folder: Optional[str] = None,
+    drop_inputs: Optional[List[str]] = None,
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Validates a model.
@@ -80,6 +145,7 @@ def validate_model(
     :param quiet: if quiet, catches exception if any issue
     :param patch: applies patches before exporting
     :param dump_folder: dumps everything in a subfolder of this one
+    :param drop_inputs: drops this list of inputs (given their names)
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -112,6 +178,27 @@ def validate_model(
     else:
         data = get_untrained_model_with_inputs(model_id, verbose=verbose, task=task)
 
+    if drop_inputs:
+        if verbose:
+            print(f"[validate_model] drop inputs {drop_inputs!r}")
+            print(f"[validate_model] current inputs: {string_type(data["inputs"])}")
+            print(
+                f"[validate_model] current dynnamic_shapes: "
+                f"{_ds_clean(data["dynamic_shapes"])}"
+            )
+        data["inputs"], data["dynamic_shapes"] = filter_inputs(
+            data["inputs"],
+            drop_names=drop_inputs,
+            model=data["model"],
+            dynamic_shapes=data["dynamic_shapes"],
+        )
+        if verbose:
+            print(f"[validate_model] new inputs: {string_type(data["inputs"])}")
+            print(
+                f"[validate_model] new dynnamic_shapes: "
+                f"{_ds_clean(data["dynamic_shapes"])}"
+            )
+
     if not empty(dtype):
         if isinstance(dtype, str):
             dtype = getattr(torch, dtype)
@@ -338,18 +425,6 @@ def call_exporter(
     )
 
 
-def split_args_kwargs(inputs: Any) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
-    """
-    Splits into args, kwargs.
-    """
-    if isinstance(inputs, dict):
-        return (), inputs
-    if isinstance(inputs, tuple) and len(inputs) == 2 and isinstance(inputs[1], dict):
-        return inputs
-    assert isinstance(inputs, tuple), f"Unexpected inputs {string_type(inputs)}"
-    return inputs, {}
-
-
 def call_torch_export_export(
     data: Dict[str, Any],
     exporter: str,