Files
ANSLibs/OpenVINO/python/openvino/utils/data_helpers/data_dispatcher.py

448 lines
15 KiB
Python

# -*- coding: utf-8 -*-
# Copyright (C) 2018-2025 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
from functools import singledispatch
from typing import Any, Union, Optional
import numpy as np
from openvino._pyopenvino import ConstOutput, Tensor, Type, RemoteTensor
from openvino.utils.data_helpers.wrappers import _InferRequestWrapper, OVDict
ContainerTypes = Union[dict, list, tuple, OVDict]
ScalarTypes = Union[np.number, int, float]
ValidKeys = Union[str, int, ConstOutput]
def is_list_simple_type(input_list: list) -> bool:
for sublist in input_list:
if isinstance(sublist, list):
for element in sublist:
if not isinstance(element, (str, float, int, bytes)):
return False
else:
if not isinstance(sublist, (str, float, int, bytes)):
return False
return True
def get_request_tensor(
request: _InferRequestWrapper,
key: Optional[ValidKeys] = None,
) -> Tensor:
if key is None:
return request.get_input_tensor()
elif isinstance(key, int):
return request.get_input_tensor(key)
elif isinstance(key, (str, ConstOutput)):
return request.get_tensor(key)
else:
raise TypeError(f"Unsupported key type: {type(key)} for Tensor under key: {key}")
@singledispatch
def value_to_tensor(
value: Union[Tensor, np.ndarray, ScalarTypes, str],
request: Optional[_InferRequestWrapper] = None,
is_shared: bool = False,
key: Optional[ValidKeys] = None,
) -> None:
raise TypeError(f"Incompatible inputs of type: {type(value)}")
@value_to_tensor.register(Tensor)
def _(
value: Tensor,
request: Optional[_InferRequestWrapper] = None,
is_shared: bool = False,
key: Optional[ValidKeys] = None,
) -> Tensor:
return value
@value_to_tensor.register(RemoteTensor)
def _(
value: RemoteTensor,
request: Optional[_InferRequestWrapper] = None,
is_shared: bool = False,
key: Optional[ValidKeys] = None,
) -> RemoteTensor:
return value
@value_to_tensor.register(np.ndarray)
def _(
value: np.ndarray,
request: _InferRequestWrapper,
is_shared: bool = False,
key: Optional[ValidKeys] = None,
) -> Tensor:
tensor = get_request_tensor(request, key)
tensor_type = tensor.get_element_type()
tensor_dtype = tensor_type.to_dtype()
# String edge-case, always copy.
# Scalars are also handled by C++.
if tensor_type == Type.string:
return Tensor(value, shared_memory=False)
# Scalars edge-case:
if value.ndim == 0:
tensor_shape = tuple(tensor.shape)
if tensor_dtype == value.dtype and tensor_shape == value.shape:
return Tensor(value, shared_memory=is_shared)
elif tensor.size == 0:
# the first infer request for dynamic input cannot reshape to 0 shape
return Tensor(value.astype(tensor_dtype).reshape((1)), shared_memory=False)
else:
return Tensor(value.astype(tensor_dtype).reshape(tensor_shape), shared_memory=False)
# WA for FP16-->BF16 edge-case, always copy.
if tensor_type == Type.bf16:
tensor = Tensor(tensor_type, value.shape)
tensor.data[:] = value.view(tensor_dtype)
return tensor
# WA for "not writeable" edge-case, always copy.
if value.flags["WRITEABLE"] is False:
tensor = Tensor(tensor_type, value.shape)
tensor.data[:] = value.astype(tensor_dtype) if tensor_dtype != value.dtype else value
return tensor
# If types are mismatched, convert and always copy.
if tensor_dtype != value.dtype:
return Tensor(value.astype(tensor_dtype), shared_memory=False)
# Otherwise, use mode defined in the call.
return Tensor(value, shared_memory=is_shared)
@value_to_tensor.register(list)
def _(
value: list,
request: _InferRequestWrapper,
is_shared: bool = False,
key: Optional[ValidKeys] = None,
) -> Tensor:
return Tensor(value)
@value_to_tensor.register(np.number)
@value_to_tensor.register(int)
@value_to_tensor.register(float)
@value_to_tensor.register(str)
@value_to_tensor.register(bytes)
def _(
value: Union[ScalarTypes, str, bytes],
request: _InferRequestWrapper,
is_shared: bool = False,
key: Optional[ValidKeys] = None,
) -> Tensor:
# np.number/int/float/str/bytes edge-case, copy will occur in both scenarios.
tensor_type = get_request_tensor(request, key).get_element_type()
tensor_dtype = tensor_type.to_dtype()
tmp = np.array(value)
# String edge-case -- it converts the data inside of Tensor class.
# If types are mismatched, convert.
if tensor_type != Type.string and tensor_dtype != tmp.dtype:
return Tensor(tmp.astype(tensor_dtype), shared_memory=False)
return Tensor(tmp, shared_memory=False)
def to_c_style(value: Any, is_shared: bool = False) -> Any:
if not isinstance(value, np.ndarray):
if hasattr(value, "__array__"):
if np.lib.NumpyVersion(np.__version__) >= "2.0.0":
# https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
return to_c_style(np.asarray(value), is_shared) if is_shared else np.asarray(value, copy=True) # type: ignore
else:
return to_c_style(np.array(value, copy=False), is_shared) if is_shared else np.array(value, copy=True)
return value
return value if value.flags["C_CONTIGUOUS"] else np.ascontiguousarray(value)
###
# Start of array normalization.
###
@singledispatch
def normalize_arrays(
inputs: Any,
is_shared: bool = False,
) -> Any:
# Check the special case of the array-interface
if hasattr(inputs, "__array__"):
if np.lib.NumpyVersion(np.__version__) >= "2.0.0":
# https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
return to_c_style(np.asarray(inputs), is_shared) if is_shared else np.asarray(inputs, copy=True) # type: ignore
else:
return to_c_style(np.array(inputs, copy=False), is_shared) if is_shared else np.array(inputs, copy=True)
# Error should be raised if type does not match any dispatchers
raise TypeError(f"Incompatible inputs of type: {type(inputs)}")
@normalize_arrays.register(dict)
def _(
inputs: dict,
is_shared: bool = False,
) -> dict:
return {k: to_c_style(v, is_shared) if is_shared else v for k, v in inputs.items()}
@normalize_arrays.register(OVDict)
def _(
inputs: OVDict,
is_shared: bool = False,
) -> dict:
return {i: to_c_style(v, is_shared) if is_shared else v for i, (_, v) in enumerate(inputs.items())}
@normalize_arrays.register(list)
@normalize_arrays.register(tuple)
def _(
inputs: Union[list, tuple],
is_shared: bool = False,
) -> dict:
return {i: to_c_style(v, is_shared) if is_shared else v for i, v in enumerate(inputs)}
@normalize_arrays.register(np.ndarray)
def _(
inputs: dict,
is_shared: bool = False,
) -> Any:
return to_c_style(inputs, is_shared) if is_shared else inputs
###
# End of array normalization.
###
###
# Start of "shared" dispatcher.
# (1) Each method should keep Tensors "as-is", regardless to them being shared or not.
# (2) ...
###
# Step to keep alive input values that are not C-style by default
@singledispatch
def create_shared(
inputs: Any,
request: _InferRequestWrapper,
) -> None:
# Check the special case of the array-interface
if hasattr(inputs, "__array__"):
request._inputs_data = normalize_arrays(inputs, is_shared=True)
return value_to_tensor(request._inputs_data, request=request, is_shared=True)
# Error should be raised if type does not match any dispatchers
raise TypeError(f"Incompatible inputs of type: {type(inputs)}")
@create_shared.register(dict)
@create_shared.register(tuple)
@create_shared.register(OVDict)
def _(
inputs: Union[dict, tuple, OVDict],
request: _InferRequestWrapper,
) -> dict:
request._inputs_data = normalize_arrays(inputs, is_shared=True)
return {k: value_to_tensor(v, request=request, is_shared=True, key=k) for k, v in request._inputs_data.items()}
# Special override to perform list-related dispatch
@create_shared.register(list)
def _(
inputs: list,
request: _InferRequestWrapper,
) -> dict:
# If list is passed to single input model and consists only of simple types
# i.e. str/bytes/float/int, wrap around it and pass into the dispatcher.
request._inputs_data = normalize_arrays([inputs] if request._is_single_input() and is_list_simple_type(inputs) else inputs, is_shared=True)
return {k: value_to_tensor(v, request=request, is_shared=True, key=k) for k, v in request._inputs_data.items()}
@create_shared.register(np.ndarray)
def _(
inputs: np.ndarray,
request: _InferRequestWrapper,
) -> Tensor:
request._inputs_data = normalize_arrays(inputs, is_shared=True)
return value_to_tensor(request._inputs_data, request=request, is_shared=True)
@create_shared.register(Tensor)
@create_shared.register(np.number)
@create_shared.register(int)
@create_shared.register(float)
@create_shared.register(str)
@create_shared.register(bytes)
def _(
inputs: Union[Tensor, ScalarTypes, str, bytes],
request: _InferRequestWrapper,
) -> Tensor:
return value_to_tensor(inputs, request=request, is_shared=True)
###
# End of "shared" dispatcher methods.
###
###
# Start of "copied" dispatcher.
###
def set_request_tensor(
request: _InferRequestWrapper,
tensor: Tensor,
key: Optional[ValidKeys] = None,
) -> None:
if key is None:
request.set_input_tensor(tensor)
elif isinstance(key, int):
request.set_input_tensor(key, tensor)
elif isinstance(key, (str, ConstOutput)):
request.set_tensor(key, tensor)
else:
raise TypeError(f"Unsupported key type: {type(key)} for Tensor under key: {key}")
@singledispatch
def update_tensor(
inputs: Any,
request: _InferRequestWrapper,
key: Optional[ValidKeys] = None,
) -> None:
if hasattr(inputs, "__array__"):
update_tensor(normalize_arrays(inputs, is_shared=False), request, key)
return None
raise TypeError(f"Incompatible inputs of type: {type(inputs)} under {key} key!")
@update_tensor.register(np.ndarray)
def _(
inputs: np.ndarray,
request: _InferRequestWrapper,
key: Optional[ValidKeys] = None,
) -> None:
if inputs.ndim != 0:
tensor = get_request_tensor(request, key)
# Update shape if there is a mismatch
if tuple(tensor.shape) != inputs.shape:
tensor.shape = inputs.shape
# When copying, type should be up/down-casted automatically.
if tensor.element_type == Type.string:
tensor.bytes_data = inputs
else:
tensor.data[:] = inputs[:]
else:
# If shape is "empty", assume this is a scalar value
set_request_tensor(
request,
value_to_tensor(inputs, request=request, is_shared=False, key=key),
key,
)
@update_tensor.register(np.number) # type: ignore
@update_tensor.register(float)
@update_tensor.register(int)
@update_tensor.register(str)
def _(
inputs: Union[ScalarTypes, str],
request: _InferRequestWrapper,
key: Optional[ValidKeys] = None,
) -> None:
set_request_tensor(
request,
value_to_tensor(inputs, request=request, is_shared=False, key=key),
key,
)
def update_inputs(inputs: dict, request: _InferRequestWrapper) -> dict:
"""Helper function to prepare inputs for inference.
It creates copy of Tensors or copy data to already allocated Tensors on device
if the item is of type `np.ndarray`, `np.number`, `int`, `float` or has numpy __array__ attribute.
If value is of type `list`, create a Tensor based on it, copy will occur in the Tensor constructor.
"""
# Create new temporary dictionary.
# new_inputs will be used to transfer data to inference calls,
# ensuring that original inputs are not overwritten with Tensors.
new_inputs: dict[ValidKeys, Tensor] = {}
for key, value in inputs.items():
if not isinstance(key, (str, int, ConstOutput)):
raise TypeError(f"Incompatible key type for input: {key}")
# Copy numpy arrays to already allocated Tensors.
# If value object has __array__ attribute, load it to Tensor using np.array
if isinstance(value, (np.ndarray, np.number, int, float, str)) or hasattr(value, "__array__"):
update_tensor(value, request, key)
elif isinstance(value, list):
new_inputs[key] = Tensor(value)
# If value is of Tensor type, put it into temporary dictionary.
elif isinstance(value, Tensor):
new_inputs[key] = value
# Throw error otherwise.
else:
raise TypeError(f"Incompatible inputs of type: {type(value)} under {key} key!")
return new_inputs
@singledispatch
def create_copied(
inputs: Union[ContainerTypes, np.ndarray, ScalarTypes, str, bytes],
request: _InferRequestWrapper,
) -> Union[dict, None]:
# Check the special case of the array-interface
if hasattr(inputs, "__array__"):
update_tensor(normalize_arrays(inputs, is_shared=False), request, key=None)
return {}
# Error should be raised if type does not match any dispatchers
raise TypeError(f"Incompatible inputs of type: {type(inputs)}")
@create_copied.register(dict)
@create_copied.register(tuple)
@create_copied.register(OVDict)
def _(
inputs: Union[dict, tuple, OVDict],
request: _InferRequestWrapper,
) -> dict:
return update_inputs(normalize_arrays(inputs, is_shared=False), request)
# Special override to perform list-related dispatch
@create_copied.register(list)
def _(
inputs: list,
request: _InferRequestWrapper,
) -> dict:
# If list is passed to single input model and consists only of simple types
# i.e. str/bytes/float/int, wrap around it and pass into the dispatcher.
return update_inputs(normalize_arrays([inputs] if request._is_single_input() and is_list_simple_type(inputs) else inputs, is_shared=False), request)
@create_copied.register(np.ndarray)
def _(
inputs: np.ndarray,
request: _InferRequestWrapper,
) -> dict:
update_tensor(normalize_arrays(inputs, is_shared=False), request, key=None)
return {}
@create_copied.register(Tensor)
@create_copied.register(np.number)
@create_copied.register(int)
@create_copied.register(float)
@create_copied.register(str)
@create_copied.register(bytes)
def _(
inputs: Union[Tensor, ScalarTypes, str, bytes],
request: _InferRequestWrapper,
) -> Tensor:
return value_to_tensor(inputs, request=request, is_shared=False)
###
# End of "copied" dispatcher methods.
###
def _data_dispatch(
request: _InferRequestWrapper,
inputs: Union[ContainerTypes, Tensor, np.ndarray, ScalarTypes, str] = None,
is_shared: bool = False,
) -> Union[dict, Tensor]:
if inputs is None:
return {}
return create_shared(inputs, request) if is_shared else create_copied(inputs, request)