openvino/tests/e2e_tests/test_utils/tf_hub_utils.py

82 lines
2.7 KiB
Python

# Copyright (C) 2018-2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import numpy as np
import tensorflow as tf
rng = np.random.default_rng(seed=56190)
type_map = {
tf.float64: np.float64,
tf.float32: np.float32,
tf.int8: np.int8,
tf.int16: np.int16,
tf.int32: np.int32,
tf.int64: np.int64,
tf.uint8: np.uint8,
tf.uint16: np.uint16,
tf.string: str,
tf.bool: bool,
}
def prepare_input(input_shape, input_type):
if input_type in [np.float32, np.float64]:
return 2.0 * rng.random(size=input_shape, dtype=input_type)
elif input_type in [np.uint8, np.uint16, np.int8, np.int16, np.int32, np.int64]:
return rng.integers(0, 5, size=input_shape).astype(input_type)
elif input_type in [str]:
return np.broadcast_to("Some string", input_shape)
elif input_type in [bool]:
return rng.integers(0, 2, size=input_shape).astype(input_type)
else:
assert False, "Unsupported type {}".format(input_type)
def prepare_inputs(inputs_info):
# if len(inputs_info) > 0 and inputs_info[0] == 'list':
# inputs = []
# inputs_info = inputs_info[1:]
# for input_name, input_shape, input_type in inputs_info:
# inputs.append(prepare_input(input_shape, input_type))
# else:
inputs = {}
for input_name, input_shape, input_type in inputs_info:
inputs[input_name] = prepare_input(input_shape, input_type)
return inputs
def get_inputs_info(model_obj):
inputs_info = []
assert len(model_obj.structured_input_signature) > 1, "incorrect model or test issue"
for input_name, input_info in model_obj.structured_input_signature[1].items():
input_shape = []
try:
if input_info.shape.as_list() == [None, None, None, 3] and input_info.dtype == tf.float32:
# image classification case, let us imitate an image
# that helps to avoid compute output size issue
input_shape = [1, 200, 200, 3]
else:
for dim in input_info.shape.as_list():
if dim is None:
input_shape.append(1)
else:
input_shape.append(dim)
except ValueError:
# unknown rank case
pass
if input_info.dtype == tf.resource:
# skip inputs corresponding to variables
continue
assert input_info.dtype in type_map, "Unsupported input type: {}".format(input_info.dtype)
inputs_info.append((input_name, input_shape, type_map[input_info.dtype]))
return inputs_info
def generate_tf_hub_inputs(model):
"""
Generates random inputs depending on model's input type
"""
return prepare_inputs(get_inputs_info(model))