support input shapes export

optimum/commands/export/openvino.py

@@ -22,6 +22,7 @@
 
 from optimum.commands.base import BaseOptimumCLICommand, CommandInfo
 from optimum.utils.constant import ALL_TASKS
+from optimum.utils import DEFAULT_DUMMY_SHAPES
 
 
 logger = logging.getLogger(__name__)
@@ -294,6 +295,82 @@ def parse_args_openvino(parser: "ArgumentParser"):
         type=json.loads,
         help=("Any kwargs passed to the model forward, or used to customize the export for a given model."),
     )
+    doc_input = "to use in the example input given to the OpenVINO export."
+    optional_group.add_argument(
+        "--batch_size",
+        type=int,
+        default=None,
+        help=f"Batch size {doc_input}",
+    )
+    optional_group.add_argument(
+        "--sequence_length",
+        type=int,
+        default=None,
+        help=f"Text tasks only. Sequence length {doc_input}",
+    )
+    optional_group.add_argument(
+        "--num_choices",
+        type=int,
+        default=None,
+        help=f"Text tasks only. Num choices {doc_input}",
+    )
+    optional_group.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        help=f"Image tasks only. Width {doc_input}",
+    )
+    optional_group.add_argument(
+        "--height",
+        type=int,
+        default=None,
+        help=f"Image tasks only. Height {doc_input}",
+    )
+    optional_group.add_argument(
+        "--num_channels",
+        type=int,
+        default=None,
+        help=f"Image tasks only. Number of channels {doc_input}",
+    )
+    optional_group.add_argument(
+        "--feature_size",
+        type=int,
+        default=None,
+        help=f"Audio tasks only. Feature size {doc_input}",
+    )
+    optional_group.add_argument(
+        "--nb_max_frames",
+        type=int,
+        default=None,
+        help=f"Audio tasks only. Maximum number of frames {doc_input}",
+    )
+    optional_group.add_argument(
+        "--audio_sequence_length",
+        type=int,
+        default=None,
+        help=f"Audio tasks only. Audio sequence length {doc_input}",
+    )
+    optional_group.add_argument(
+        "--point_batch_size",
+        type=int,
+        default=None,
+        help=(
+            "For Segment Anything. It corresponds to how many segmentation masks we want the model to predict per "
+            "input point."
+        ),
+    )
+    optional_group.add_argument(
+        "--nb_points_per_image",
+        type=int,
+        default=None,
+        help="For Segment Anything. It corresponds to the number of points per segmentation masks.",
+    )
+    optional_group.add_argument(
+        "--visual_seq_length",
+        type=int,
+        default=None,
+        help="Visual sequence length",
+    )
 
 
 def no_compression_parameter_provided(args):
@@ -464,7 +541,11 @@ def run(self):
             output = Path(self.args.output)
 
         try:
-            # TODO : add input shapes
+            # Get the shapes to be used to generate dummy inputs
+            input_shapes = {}
+            for input_name in DEFAULT_DUMMY_SHAPES:
+                if hasattr(self.args, input_name) and getattr(self.args, input_name) is not None:
+                    input_shapes[input_name] = getattr(self.args, input_name)
             main_export(
                 model_name_or_path=self.args.model,
                 output=output,
@@ -479,7 +560,7 @@ def run(self):
                 library_name=library_name,
                 variant=self.args.variant,
                 model_kwargs=self.args.model_kwargs,
-                # **input_shapes,
+                **input_shapes,
             )
             if apply_main_quantize:
                 _main_quantize(

optimum/exporters/openvino/convert.py

@@ -361,8 +361,10 @@ 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
+        static_shapes = True
+        if input_shapes["static_shapes"] == False:
+            input_shapes = {}  # will use dynamic shapes in the OpenVINO IR
+            static_shapes = False  # User did not explicit input shapes, produce dynamic shapes in the model
 
         # Check that inputs match, and order them properly
         dummy_inputs = config.generate_dummy_inputs(framework="pt", **input_shapes)
@@ -411,7 +413,7 @@ def ts_patched_forward(*args, **kwargs):
 
         with patcher:
             check_dummy_inputs_are_allowed(model, dummy_inputs)
-            input_info = _get_input_info(model, config, dummy_inputs)
+            input_info = _get_input_info(model, config, dummy_inputs, static_shapes=static_shapes)
             torch_export = os.getenv("OPENVINO_DYNAMO_EXPORT", "false").lower() == "true"
             if torch_export:
                 if hasattr(torch.ops, "_prepare_4d_causal_attention_mask_for_sdpa"):
@@ -633,12 +635,13 @@ def export_from_model(
     # Get the shapes to be used to generate dummy inputs
     input_shapes = {}
     for input_name in DEFAULT_DUMMY_SHAPES.keys():
-        if input_name in ["height", "width"]:
+        if input_name in ["height", "width"] and kwargs_shapes is None:
             # use H and W from generator defaults
             continue
         input_shapes[input_name] = (
             kwargs_shapes[input_name] if input_name in kwargs_shapes else DEFAULT_DUMMY_SHAPES[input_name]
         )
+    input_shapes["static_shapes"] = True if kwargs_shapes else False
 
     if library_name == "open_clip":
         custom_architecture = True

optimum/exporters/openvino/utils.py

@@ -87,7 +87,10 @@ def flattenize_inputs(inputs: List[Any]):
 
 
 def _get_input_info(
-    model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, dummy_inputs: Dict[str, Any]
+    model: Union["PreTrainedModel", "ModelMixin"],
+    config: OnnxConfig,
+    dummy_inputs: Dict[str, Any],
+    static_shapes: bool = False,
 ) -> List[InputInfo]:
     sig = inspect.signature(model.forward) if hasattr(model, "forward") else inspect.signature(model.call)
     inputs = config.ordered_inputs(model)
@@ -105,7 +108,8 @@ def _get_input_info(
         example = flatten_inputs[i]
         type = get_element_type(example.cpu().numpy().dtype)
         shape = PartialShape(example.shape)
-        if name in inputs:
+        # Set dynamic axes from config, unless user requested static shapes
+        if name in inputs and not static_shapes:
             named_dims = inputs[name]
             for idx, dim_name in named_dims.items():
                 if dim_name in name_to_symbol: