[TF FE] Refactor Unpack and add layer test (#15519)
* [TF FE] Refactor Unpack and add layer test Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com> * Update tests/layer_tests/tensorflow_tests/test_tf_Unpack.py --------- Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>
This commit is contained in:
parent
46564543f8
commit
b544308616
@ -14,23 +14,20 @@ namespace tensorflow {
|
||||
namespace op {
|
||||
|
||||
OutputVector translate_unpack_op(const NodeContext& node) {
|
||||
TENSORFLOW_OP_VALIDATION(node, node.get_input_size() > 0, "Unpack must have at least one input.");
|
||||
auto input = node.get_input(0);
|
||||
default_op_checks(node, 1, {"Unpack", "UNPACK"});
|
||||
auto value = node.get_input(0);
|
||||
auto axis = node.get_attribute<int64_t>("axis", 0);
|
||||
auto num = node.get_attribute<int64_t>("num");
|
||||
|
||||
auto axis_const = make_shared<Constant>(element::i64, Shape{}, axis);
|
||||
auto split = make_shared<Split>(input, axis_const, num);
|
||||
OutputVector res;
|
||||
int idx = 0;
|
||||
for (auto out : split->outputs()) {
|
||||
auto squeezed_res = make_shared<Squeeze>(out, axis_const);
|
||||
squeezed_res->set_friendly_name(node.get_name() + "/squeeze_" + to_string(idx));
|
||||
set_out_name(node.get_name() + ":" + std::to_string(idx), squeezed_res->output(0));
|
||||
++idx;
|
||||
res.push_back(squeezed_res);
|
||||
auto split = make_shared<Split>(value, axis_const, num);
|
||||
OutputVector unpack_outputs;
|
||||
for (int output_ind = 0; output_ind < num; ++output_ind) {
|
||||
auto unpack_output = make_shared<Squeeze>(split->output(output_ind), axis_const);
|
||||
set_out_name(node.get_name() + ":" + to_string(output_ind), unpack_output);
|
||||
unpack_outputs.push_back(unpack_output);
|
||||
}
|
||||
return res;
|
||||
return unpack_outputs;
|
||||
}
|
||||
} // namespace op
|
||||
} // namespace tensorflow
|
||||
|
55
tests/layer_tests/tensorflow_tests/test_tf_Unpack.py
Normal file
55
tests/layer_tests/tensorflow_tests/test_tf_Unpack.py
Normal file
@ -0,0 +1,55 @@
|
||||
# Copyright (C) 2018-2023 Intel Corporation
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import tensorflow as tf
|
||||
from common.tf_layer_test_class import CommonTFLayerTest
|
||||
|
||||
|
||||
class TestUnpack(CommonTFLayerTest):
|
||||
def _prepare_input(self, inputs_info):
|
||||
assert 'x' in inputs_info, "Test error: inputs_info must contain `x`"
|
||||
x_shape = inputs_info['x']
|
||||
inputs_data = {}
|
||||
inputs_data['x'] = np.random.randint(-10, 10, x_shape).astype(self.input_type)
|
||||
return inputs_data
|
||||
|
||||
def create_unpack_net(self, input_shape, num, axis, input_type):
|
||||
self.input_type = input_type
|
||||
tf.compat.v1.reset_default_graph()
|
||||
# Create the graph and model
|
||||
with tf.compat.v1.Session() as sess:
|
||||
type_map = {
|
||||
np.float32: tf.float32,
|
||||
np.int32: tf.int32,
|
||||
}
|
||||
assert input_type in type_map, "Test error: need to update type_map"
|
||||
tf_type = type_map[input_type]
|
||||
x = tf.compat.v1.placeholder(tf_type, input_shape, 'x')
|
||||
if axis is not None:
|
||||
unpack = tf.raw_ops.Unpack(value=x, num=num, axis=axis)
|
||||
else:
|
||||
unpack = tf.raw_ops.Unpack(value=x, num=num)
|
||||
for ind in range(num):
|
||||
tf.identity(unpack[ind], name="output_" + str(ind))
|
||||
tf.compat.v1.global_variables_initializer()
|
||||
|
||||
tf_net = sess.graph_def
|
||||
|
||||
return tf_net, None
|
||||
|
||||
test_data_basic = [
|
||||
dict(input_shape=[2, 3, 4], num=3, axis=1, input_type=np.float32),
|
||||
dict(input_shape=[3, 4], num=3, axis=None, input_type=np.int32),
|
||||
dict(input_shape=[4, 2, 3], num=2, axis=-2, input_type=np.float32),
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("params", test_data_basic)
|
||||
@pytest.mark.precommit_tf_fe
|
||||
@pytest.mark.nightly
|
||||
def test_unpack_basic(self, params, ie_device, precision, ir_version, temp_dir,
|
||||
use_new_frontend, use_old_api):
|
||||
self._test(*self.create_unpack_net(**params),
|
||||
ie_device, precision, ir_version, temp_dir=temp_dir,
|
||||
use_new_frontend=use_new_frontend, use_old_api=use_old_api)
|
Loading…
Reference in New Issue
Block a user