[OV] Sets symbols on inputs and variables (#696)

* Prepare shape with symbols before conversion

* Bumping lower bound of openvino version due to backward incompatible optimum change

* style

* style

Co-authored-by: Ekaterina Aidova <ekaterina.aidova@intel.com>

* style

* style

* Style

* rename function

* fix test

---------

Co-authored-by: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com>
Co-authored-by: Ekaterina Aidova <ekaterina.aidova@intel.com>
Co-authored-by: Ella Charlaix <ella@huggingface.co>

Author: 4 people

optimum/exporters/openvino/convert.py

@@ -14,7 +14,6 @@
 
 import functools
 import gc
-import inspect
 import logging
 import os
 from pathlib import Path
@@ -23,9 +22,8 @@
 import onnx
 from transformers.utils import is_tf_available, is_torch_available
 
-from openvino.runtime import Model, PartialShape, save_model
+from openvino.runtime import Model, save_model
 from openvino.runtime.exceptions import OVTypeError
-from openvino.runtime.utils.types import get_element_type
 from openvino.tools.ovc import convert_model
 from optimum.exporters import TasksManager
 from optimum.exporters.onnx.base import OnnxConfig
@@ -50,9 +48,8 @@
 from .stateful import ensure_export_task_support_stateful, ensure_stateful_is_available, patch_stateful
 from .utils import (
     OV_XML_FILE_NAME,
+    _get_input_info,
     clear_class_registry,
-    flattenize_inputs,
-    get_input_shapes,
     remove_none_from_dummy_inputs,
 )
 
@@ -374,13 +371,12 @@ def ts_patched_forward(*args, **kwargs):
 
                     __make_16bit_traceable(model)
                 check_dummy_inputs_are_allowed(model, dummy_inputs)
-                sig = inspect.signature(model.forward) if hasattr(model, "forward") else inspect.signature(model.call)
-                inputs = config.ordered_inputs(model)
-                input_names = list(inputs.keys())
-                output_names = list(config.outputs.keys())
-                input_info = get_input_shapes(dummy_inputs, inputs)
-
-                ov_model = convert_model(model, example_input=dummy_inputs, input=input_info)
+                input_info = _get_input_info(model, config, dummy_inputs)
+                ov_model = convert_model(
+                    model,
+                    example_input=dummy_inputs,
+                    input=[(item.shape, item.type) for item in input_info],
+                )
 
         except Exception as ex:
             logger.warning(f"Export model to OpenVINO directly failed with: \n{ex}.\nModel will be exported to ONNX")
@@ -411,27 +407,17 @@ def ts_patched_forward(*args, **kwargs):
                 ov_config=ov_config,
             )
 
-        ordered_dummy_inputs = {param: dummy_inputs[param] for param in sig.parameters if param in dummy_inputs}
-        if not ordered_dummy_inputs:
-            ordered_dummy_inputs = dummy_inputs
-        ordered_input_names = list(inputs)
-        flatten_inputs = flattenize_inputs(ordered_dummy_inputs.values())
-        ov_model.validate_nodes_and_infer_types()
+        ov_model.validate_nodes_and_infer_types()  # TODO: remove as unnecessary validation?
+
+        output_names = list(config.outputs.keys())
         for idx, out_tensor in enumerate(ov_model.outputs):
             if idx < len(output_names):
                 out_tensor.get_tensor().set_names({output_names[idx]})
 
+        input_names = [item.name for item in input_info]
         for idx, inp_tensor in enumerate(ov_model.inputs):
-            input_name = ordered_input_names[idx]
+            input_name = input_names[idx]
             inp_tensor.get_tensor().set_names({input_name})
-            inp_data = flatten_inputs[idx]
-            static_shape = PartialShape(inp_data.shape)
-            dims = inputs.get(input_name, [])
-            for dim in dims:
-                static_shape[dim] = -1
-            inp_tensor.get_node().set_partial_shape(static_shape)
-            inp_tensor.get_node().set_element_type(get_element_type(inp_data.cpu().numpy().dtype))
-        ov_model.validate_nodes_and_infer_types()
 
         if stateful:
             patch_stateful(model.config, ov_model)

optimum/exporters/openvino/utils.py

@@ -12,14 +12,21 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import inspect
+from collections import namedtuple
 from typing import Any, Dict, List, Tuple, Union
 
 from transformers.utils import is_torch_available
 
-from openvino.runtime import PartialShape
+from openvino.runtime import Dimension, PartialShape, Symbol
+from openvino.runtime.utils.types import get_element_type
+from optimum.exporters.onnx.base import OnnxConfig
 from optimum.utils import is_diffusers_available
 
 
+InputInfo = namedtuple("InputInfo", ["name", "shape", "type", "example"])
+
+
 if is_torch_available():
     import torch
     import torch.nn as nn
@@ -69,6 +76,41 @@ def flattenize_inputs(inputs: List[Any]):
     return flatten_inputs
 
 
+def _get_input_info(
+    model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, dummy_inputs: Dict[str, Any]
+) -> List[InputInfo]:
+    sig = inspect.signature(model.forward) if hasattr(model, "forward") else inspect.signature(model.call)
+    inputs = config.ordered_inputs(model)
+    ordered_dummy_inputs = {param: dummy_inputs[param] for param in sig.parameters if param in dummy_inputs}
+    if not ordered_dummy_inputs:
+        ordered_dummy_inputs = dummy_inputs
+    ordered_input_names = list(inputs)
+    flatten_inputs = flattenize_inputs(ordered_dummy_inputs.values())
+    input_info = []
+
+    name_to_symbol = {}
+
+    for i in range(len(ordered_input_names)):
+        name = ordered_input_names[i]
+        example = flatten_inputs[i]
+        type = get_element_type(example.cpu().numpy().dtype)
+        shape = PartialShape(example.shape)
+        if name in inputs:
+            named_dims = inputs[name]
+            for idx, dim_name in named_dims.items():
+                if dim_name in name_to_symbol:
+                    symbol = name_to_symbol[dim_name]
+                else:
+                    symbol = Symbol()
+                    name_to_symbol[name] = symbol
+                dim = Dimension(-1)
+                dim.set_symbol(symbol)
+                shape[idx] = dim
+        info = InputInfo(name=name, shape=shape, type=type, example=example)
+        input_info.append(info)
+    return input_info
+
+
 def remove_none_from_dummy_inputs(dummy_inputs: Dict[str, Any]):
     """
     Removes None values from the dictionary.
@@ -109,31 +151,6 @@ def remove_none_from_list_tuple(item: Union[List[Any], Tuple[Any]]):
     return upd_dummy, dict_dummy
 
 
-def get_input_shapes(dummy_inputs: Dict[str, Any], inputs: Dict[str, Any]):
-    """
-    Resolves input shapes based on dynamic axes from input config and dummy input shapes
-
-    Args:
-        dummy_inputs (Dict[str, Any]): A dictionary of dummy inputs.
-        inputs (Dict[str, Any]): A dictionary of input tensors.
-
-    Returns:
-       input_info: List of input info for conversion
-
-    """
-    input_info = []
-    for input_name, data in dummy_inputs.items():
-        if isinstance(data, (tuple, list, dict)):
-            return None
-        static_shape = PartialShape(data.shape)
-        if input_name in inputs:
-            dynamic_dims = inputs[input_name]
-            for dim in dynamic_dims:
-                static_shape[dim] = -1
-        input_info.append((input_name, static_shape))
-    return input_info
-
-
 def clear_class_registry():
     """
     Removes Torchscript cached modules