Support for custom input shapes in exporters onnx (#575)

* add support for custom input shapes in exporters onnx

* add tests

* fix test

* fix bug

* Update tests/exporters/test_onnx_export.py

Co-authored-by: regisss <15324346+regisss@users.noreply.github.com>

* remove run slow

* fix _get_models_to_test

* fixup post merge

* yet fixup post merge

* use f string

* authorize unused arguments in input generators

* use hf-internal-testing models

* remove redundant argument

* fixup merge

* remove redundant arg

* fix messed merge

* change blenderbot

* workign deit

* valid gptj

* valid marian

* add whisper

* fix test

* fix marian

* nit

* fix bloom

Co-authored-by: regisss <15324346+regisss@users.noreply.github.com>

.github/workflows/test_exporters.yml

@@ -31,4 +31,4 @@ jobs:
     - name: Test with unittest
       working-directory: tests
       run: |
-        RUN_SLOW=1 pytest exporters --durations=0
+        pytest exporters -s --durations=0

optimum/commands/export/onnx.py

@@ -15,21 +15,49 @@
 import subprocess
 from pathlib import Path
 
+from ...exporters import TasksManager
+from ...utils import DEFAULT_DUMMY_SHAPES
+
 
 def parse_args_onnx(parser):
-    parser.add_argument(
+    required_group = parser.add_argument_group("Required arguments")
+    required_group.add_argument(
         "-m", "--model", type=str, required=True, help="Model ID on huggingface.co or path on disk to load model from."
     )
-    parser.add_argument(
+    required_group.add_argument(
+        "output", type=Path, help="Path indicating the directory where to store generated ONNX model."
+    )
+
+    optional_group = parser.add_argument_group("Optional arguments")
+    optional_group.add_argument(
         "--task",
         default="auto",
-        help="The type of task to export the model with.",
+        help=(
+            "The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:"
+            f" {str(list(TasksManager._TASKS_TO_AUTOMODELS.keys()))}. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder."
+        ),
+    )
+    optional_group.add_argument(
+        "--for-ort",
+        action="store_true",
+        help=(
+            "This exports models ready to be run with Optimum's ORTModel. Useful for encoder-decoder models for"
+            "conditional generation. If enabled the encoder and decoder of the model are exported separately."
+        ),
     )
-    parser.add_argument("--opset", type=int, default=None, help="ONNX opset version to export the model with.")
-    parser.add_argument(
-        "--atol", type=float, default=None, help="Absolute difference tolerance when validating the model."
+    optional_group.add_argument(
+        "--opset",
+        type=int,
+        default=None,
+        help="If specified, ONNX opset version to export the model with. Otherwise, the default opset will be used.",
+    )
+    optional_group.add_argument(
+        "--atol",
+        type=float,
+        default=None,
+        help="If specified, the absolute difference tolerance when validating the model. Otherwise, the default atol for the model will be used.",
     )
-    parser.add_argument(
+    optional_group.add_argument(
         "--framework",
         type=str,
         choices=["pt", "tf"],
@@ -40,7 +68,7 @@ def parse_args_onnx(parser):
             " or what is available in the environment."
         ),
     )
-    parser.add_argument(
+    optional_group.add_argument(
         "--pad_token_id",
         type=int,
         default=None,
@@ -49,16 +77,66 @@ def parse_args_onnx(parser):
             " it."
         ),
     )
-    parser.add_argument("--cache_dir", type=str, default=None, help="Path indicating where to store cache.")
-    parser.add_argument(
-        "--for-ort",
-        action="store_true",
-        help=(
-            "This exports models ready to be run with optimum.onnxruntime. Useful for encoder-decoder models for"
-            "conditional generation. If enabled the encoder and decoder of the model are exported separately."
-        ),
+    optional_group.add_argument("--cache_dir", type=str, default=None, help="Path indicating where to store cache.")
+
+    input_group = parser.add_argument_group(
+        "Input shapes (if necessary, this allows to override the shapes of the input given to the ONNX exporter, that requires an example input.)"
+    )
+    doc_input = "to use in the example input given to the ONNX export."
+    input_group.add_argument(
+        "--batch_size",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["batch_size"],
+        help=f"Text tasks only. Batch size {doc_input}",
+    )
+    input_group.add_argument(
+        "--sequence_length",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["sequence_length"],
+        help=f"Text tasks only. Sequence length {doc_input}",
+    )
+    input_group.add_argument(
+        "--num_choices",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["num_choices"],
+        help=f"Text tasks only. Num choices {doc_input}",
+    )
+    input_group.add_argument(
+        "--width",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["width"],
+        help=f"Image tasks only. Width {doc_input}",
+    )
+    input_group.add_argument(
+        "--height",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["height"],
+        help=f"Image tasks only. Height {doc_input}",
+    )
+    input_group.add_argument(
+        "--num_channels",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["num_channels"],
+        help=f"Image tasks only. Number of channels {doc_input}",
+    )
+    input_group.add_argument(
+        "--feature_size",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["feature_size"],
+        help=f"Audio tasks only. Feature size {doc_input}",
+    )
+    input_group.add_argument(
+        "--nb_max_frames",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["nb_max_frames"],
+        help=f"Audio tasks only. Maximum number of frames {doc_input}",
+    )
+    input_group.add_argument(
+        "--audio_sequence_length",
+        type=int,
+        default=DEFAULT_DUMMY_SHAPES["audio_sequence_length"],
+        help=f"Audio tasks only. Audio sequence length {doc_input}",
     )
-    parser.add_argument("output", type=Path, help="Path indicating the directory where to store generated ONNX model.")
 
 
 class ONNXExportCommand:

optimum/exporters/onnx/__main__.py

@@ -20,7 +20,7 @@
 from transformers import AutoTokenizer
 
 from ...commands.export.onnx import parse_args_onnx
-from ...utils import is_diffusers_available, logging
+from ...utils import DEFAULT_DUMMY_SHAPES, logging
 from ...utils.save_utils import maybe_save_preprocessors
 from ..tasks import TasksManager
 from .base import OnnxConfigWithPast
@@ -66,11 +66,15 @@ def main():
                 f"The task could not be automatically inferred. Please provide the argument --task with the task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
             )
 
-    # Allocate the model
+    # get the shapes to be used to generate dummy inputs
+    input_shapes = {}
+    for input_name in DEFAULT_DUMMY_SHAPES.keys():
+        input_shapes[input_name] = getattr(args, input_name)
+
     model = TasksManager.get_model_from_task(task, args.model, framework=args.framework, cache_dir=args.cache_dir)
 
     if task != "stable-diffusion":
-        onnx_config_constructor = TasksManager.get_exporter_config_constructor(model, "onnx", task=task)
+        onnx_config_constructor = TasksManager.get_exporter_config_constructor(model=model, exporter="onnx", task=task)
         onnx_config = onnx_config_constructor(model.config)
 
         needs_pad_token_id = (
@@ -134,9 +138,12 @@ def main():
             opset=args.opset,
             output_dir=args.output.parent,
             output_names=output_names,
+            input_shapes=input_shapes,
         )
     else:
-        onnx_inputs, onnx_outputs = export(model=model, config=onnx_config, output=args.output, opset=args.opset)
+        onnx_inputs, onnx_outputs = export(
+            model=model, config=onnx_config, output=args.output, opset=args.opset, input_shapes=input_shapes
+        )
 
     try:
         if task == "stable-diffusion" or (

optimum/exporters/onnx/base.py

@@ -189,8 +189,7 @@ def _create_dummy_input_generator_classes(self, **kwargs) -> List[DummyInputGene
         """
         first_inputs_gen = self.DUMMY_INPUT_GENERATOR_CLASSES[0](self.task, self._normalized_config, **kwargs)
         dummy_inputs_generators = [
-            cls_(self.task, self._normalized_config, batch_size=first_inputs_gen.batch_size, **kwargs)
-            for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES[1:]
+            cls_(self.task, self._normalized_config, **kwargs) for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES[1:]
         ]
         dummy_inputs_generators.insert(0, first_inputs_gen)
 
@@ -274,6 +273,7 @@ def ordered_inputs(self, model: Union["PreTrainedModel", "TFPreTrainedModel"]) -
             sig = inspect.signature(model.forward)
         else:
             sig = inspect.signature(model.call)
+
         for param in sig.parameters:
             param_regex = re.compile(rf"{param}(\.\d*)?")
             to_insert = []

optimum/exporters/onnx/config.py

@@ -115,13 +115,11 @@ def _create_dummy_input_generator_classes(self, **kwargs) -> List["DummyInputGen
         dummy_decoder_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[1](
             self.task,
             self._normalized_config,
-            batch_size=dummy_text_input_generator.batch_size,
             **kwargs,
         )
         dummy_seq2seq_past_key_values_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[2](
             self.task,
             self._normalized_config,
-            batch_size=dummy_text_input_generator.batch_size,
             encoder_sequence_length=dummy_text_input_generator.sequence_length,
             **kwargs,
         )

optimum/exporters/onnx/convert.py

@@ -85,6 +85,7 @@ def validate_models_outputs(
     output_dir: Path,
     atol: Optional[float] = None,
     output_names: Optional[List[str]] = None,
+    input_shapes: Optional[Dict] = None,
 ):
     """
     Validates the export of several models, by checking that the outputs from both the reference and the exported model match.
@@ -102,10 +103,11 @@ def validate_models_outputs(
         output_names (`Optional[List[str]]`, defaults to `None`):
             The names to use for the exported ONNX files. The order must be the same as the order of submodels in the ordered dict `models_and_onnx_configs`.
             If None, will use the keys from the `models_and_onnx_configs` as names.
+        input_shapes (`Optional[Dict]`, defaults to `None`):
+            If specified, allows to use specific shapes to validate the ONNX model on.
     Raises:
         ValueError: If the outputs shapes or values do not match between the reference and the exported model.
     """
-
     if len(onnx_named_outputs) != len(models_and_onnx_configs.keys()):
         raise ValueError(
             f"Invalid number of ONNX named outputs. Required {len(models_and_onnx_configs.keys())}, Provided {len(onnx_named_outputs)}"
@@ -129,6 +131,7 @@ def validate_models_outputs(
             onnx_model=onnx_model_path,
             onnx_named_outputs=onnx_named_outputs[i],
             atol=atol,
+            input_shapes=input_shapes,
         )
 
 
@@ -138,6 +141,7 @@ def validate_model_outputs(
     onnx_model: Path,
     onnx_named_outputs: List[str],
     atol: Optional[float] = None,
+    input_shapes: Optional[Dict] = None,
 ):
     """
     Validates the export by checking that the outputs from both the reference and the exported model match.
@@ -153,6 +157,8 @@ def validate_model_outputs(
             The names of the outputs to check.
         atol (`Optional[float]`, defaults to `None`):
             The absolute tolerance in terms of outputs difference between the reference and the exported model.
+        input_shapes (`Optional[Dict]`, defaults to `None`):
+            If specified, allows to use specific shapes to validate the ONNX model on.
 
     Raises:
         ValueError: If the outputs shapes or values do not match between the reference and the exported model.
@@ -168,7 +174,10 @@ def validate_model_outputs(
         raise ImportError("The pip package `diffusers` is required to validate stable diffusion ONNX models.")
 
     framework = "pt" if is_torch_available() and isinstance(reference_model, nn.Module) else "tf"
-    reference_model_inputs = config.generate_dummy_inputs(framework=framework)
+
+    if input_shapes is None:
+        input_shapes = {}  # will use the defaults from DEFAULT_DUMMY_SHAPES
+    reference_model_inputs = config.generate_dummy_inputs(framework=framework, **input_shapes)
 
     # Create ONNX Runtime session
     options = SessionOptions()
@@ -269,6 +278,7 @@ def export_pytorch(
     opset: int,
     output: Path,
     device: str = "cpu",
+    input_shapes: Optional[Dict] = None,
 ) -> Tuple[List[str], List[str]]:
     """
     Exports a PyTorch model to an ONNX Intermediate Representation.
@@ -285,6 +295,8 @@ def export_pytorch(
         device (`str`, *optional*, defaults to `cpu`):
             The device on which the ONNX model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for
             export on CUDA devices.
+        input_shapes (`optional[Dict]`, defaults to `None`):
+            If specified, allows to use specific shapes for the example input provided to the ONNX exporter.
 
     Returns:
         `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
@@ -307,8 +319,11 @@ def export_pytorch(
                 logger.info(f"\t- {override_config_key} -> {override_config_value}")
                 setattr(model.config, override_config_key, override_config_value)
 
+        if input_shapes is None:
+            input_shapes = {}  # will use the defaults from DEFAULT_DUMMY_SHAPES
+
         # Check that inputs match, and order them properly
-        dummy_inputs = config.generate_dummy_inputs(framework="pt")
+        dummy_inputs = config.generate_dummy_inputs(framework="pt", **input_shapes)
         device = torch.device(device)
         if device.type == "cuda" and torch.cuda.is_available():
             model.to(device)
@@ -426,6 +441,7 @@ def export_models(
     opset: Optional[int] = None,
     output_names: Optional[List[str]] = None,
     device: str = "cpu",
+    input_shapes: Optional[Dict] = None,
 ) -> Tuple[List[List[str]], List[List[str]]]:
     """
     Exports a Pytorch or TensorFlow encoder decoder model to an ONNX Intermediate Representation.
@@ -434,7 +450,7 @@ def export_models(
 
     Args:
         models_and_onnx_configs (`Dict[str, Tuple[Union[`PreTrainedModel`, `TFPreTrainedModel`], `OnnxConfig`]]):
-             A dictionnary containing the models to export and their corresponding onnx configs.
+            A dictionnary containing the models to export and their corresponding onnx configs.
         output_dir (`Path`):
             Output directory to store the exported ONNX models.
         opset (`Optional[int]`, defaults to `None`):
@@ -445,6 +461,8 @@ def export_models(
         device (`str`, *optional*, defaults to `cpu`):
             The device on which the ONNX model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for
             export on CUDA devices.
+        input_shapes (`Optional[Dict]`, defaults to `None`):
+            If specified, allows to use specific shapes for the example input provided to the ONNX exporter.
     Returns:
         `Tuple[List[List[str]], List[List[str]]]`: A tuple with an ordered list of the model's inputs, and the named
         inputs from the ONNX configuration.
@@ -470,6 +488,7 @@ def export_models(
                 output=output_path,
                 opset=opset,
                 device=device,
+                input_shapes=input_shapes,
             )
         )
 
@@ -483,6 +502,7 @@ def export(
     output: Path,
     opset: Optional[int] = None,
     device: str = "cpu",
+    input_shapes: Optional[Dict] = None,
 ) -> Tuple[List[str], List[str]]:
     """
     Exports a Pytorch or TensorFlow model to an ONNX Intermediate Representation.
@@ -499,6 +519,8 @@ def export(
         device (`str`, *optional*, defaults to `cpu`):
             The device on which the ONNX model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for
             export on CUDA devices.
+        input_shapes (`Optional[Dict]`, defaults to `None`):
+            If specified, allows to use specific shapes for the example input provided to the ONNX exporter.
 
     Returns:
         `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
@@ -531,11 +553,13 @@ def export(
                 f"Unsupported PyTorch version for this model. Minimum required is {config.MIN_TORCH_VERSION},"
                 f" got: {torch.__version__}"
             )
-        return export_pytorch(model, config, opset, output, device=device)
+        return export_pytorch(model, config, opset, output, device=device, input_shapes=input_shapes)
 
     elif is_tf_available() and issubclass(type(model), TFPreTrainedModel):
         if device == "cuda":
             raise RuntimeError("`tf2onnx` does not support export on CUDA device.")
+        if input_shapes is not None:
+            logger.info("`input_shapes` argument is not supported by the Tensorflow ONNX export and will be ignored.")
         return export_tensorflow(model, config, opset, output)
 
     else: