diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py index cd3280189e..012e813283 100644 --- a/optimum/commands/export/openvino.py +++ b/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( diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index 0ed08a5f2f..88ca40fd1c 100644 --- a/optimum/exporters/openvino/convert.py +++ b/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 diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index eea8af6b6d..b11d9ee9a4 100644 --- a/optimum/exporters/openvino/utils.py +++ b/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: