From 4e175418f2224ef7d20673aaed81195d6039d95f Mon Sep 17 00:00:00 2001 From: Anastasiia Pnevskaia Date: Fri, 13 Sep 2024 22:25:23 +0200 Subject: [PATCH] [TF FE] SparseSegmentMean translator (#26540) ### Details: - Added translator for SparseSegmentMean ### Tickets: - 149705 --------- Co-authored-by: Roman Kazantsev --- .../tensorflow/docs/supported_ops.md | 2 +- src/frontends/tensorflow/src/op_table.cpp | 1 + .../include/common_op_table.hpp | 1 + .../src/op/sparse_segment_mean.cpp | 75 ++++++++++++++++++ .../test_tf_SparseSegmentMean.py | 79 +++++++++++++++++++ 5 files changed, 157 insertions(+), 1 deletion(-) create mode 100644 src/frontends/tensorflow_common/src/op/sparse_segment_mean.cpp create mode 100644 tests/layer_tests/tensorflow_tests/test_tf_SparseSegmentMean.py diff --git a/src/frontends/tensorflow/docs/supported_ops.md b/src/frontends/tensorflow/docs/supported_ops.md index 1affaa2640cbea..31f871aca4fdef 100644 --- a/src/frontends/tensorflow/docs/supported_ops.md +++ b/src/frontends/tensorflow/docs/supported_ops.md @@ -1130,7 +1130,7 @@ A "supported operation" is one that TensorFlow Frontend can convert to the OpenV | SparseReduceSumSparse | NO | | | SparseReorder | NO | | | SparseReshape | YES | | -| SparseSegmentMean | NO | | +| SparseSegmentMean | YES | | | SparseSegmentMeanGrad | NO | | | SparseSegmentMeanGradV2 | NO | | | SparseSegmentMeanWithNumSegments | NO | | diff --git a/src/frontends/tensorflow/src/op_table.cpp b/src/frontends/tensorflow/src/op_table.cpp index 0a1baf3034df42..c2bf460b02b19a 100644 --- a/src/frontends/tensorflow/src/op_table.cpp +++ b/src/frontends/tensorflow/src/op_table.cpp @@ -379,6 +379,7 @@ const std::map get_supported_ops() { {"Softmax", CreatorFunction(translate_softmax_op)}, {"SpaceToDepth", CreatorFunction(translate_space_to_depth_op)}, {"SparseReshape", CreatorFunction(translate_sparse_reshape_op)}, + {"SparseSegmentMean", CreatorFunction(translate_sparse_segment_mean_op)}, {"SparseTensorDenseAdd", CreatorFunction(translate_sparse_tensor_dense_add_op)}, {"SparseTensorDenseMatMul", CreatorFunction(translate_sparse_tensor_dense_mat_mul_op)}, {"SparseToDense", CreatorFunction(translate_sparse_to_dense_op)}, diff --git a/src/frontends/tensorflow_common/include/common_op_table.hpp b/src/frontends/tensorflow_common/include/common_op_table.hpp index a59a5db7dfa080..1fa3386fce47af 100644 --- a/src/frontends/tensorflow_common/include/common_op_table.hpp +++ b/src/frontends/tensorflow_common/include/common_op_table.hpp @@ -148,6 +148,7 @@ OP_CONVERTER(translate_rsqrt_op); OP_CONVERTER(translate_scatter_nd_op); OP_CONVERTER(translate_segment_sum_op); OP_CONVERTER(translate_space_to_batch_nd_op); +OP_CONVERTER(translate_sparse_segment_mean_op); OP_CONVERTER(translate_sparse_tensor_dense_add_op); OP_CONVERTER(translate_sparse_tensor_dense_mat_mul_op); OP_CONVERTER(translate_sparse_to_dense_op); diff --git a/src/frontends/tensorflow_common/src/op/sparse_segment_mean.cpp b/src/frontends/tensorflow_common/src/op/sparse_segment_mean.cpp new file mode 100644 index 00000000000000..d275f3a3460ce5 --- /dev/null +++ b/src/frontends/tensorflow_common/src/op/sparse_segment_mean.cpp @@ -0,0 +1,75 @@ +// Copyright (C) 2018-2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "common_op_table.hpp" +#include "openvino/op/add.hpp" +#include "openvino/op/broadcast.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/convert.hpp" +#include "openvino/op/divide.hpp" +#include "openvino/op/embedding_segments_sum.hpp" +#include "openvino/op/gather.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/scatter_update.hpp" +#include "openvino/op/shape_of.hpp" +#include "openvino/op/squeeze.hpp" +#include "openvino/op/subtract.hpp" +#include "openvino/op/transpose.hpp" +#include "openvino/op/unique.hpp" +#include "utils.hpp" + +using namespace std; +using namespace ov::op; + +namespace ov { +namespace frontend { +namespace tensorflow { +namespace op { +OutputVector translate_sparse_segment_mean_op(const NodeContext& node) { + default_op_checks(node, 3, {"SparseSegmentMean"}); + auto data = node.get_input(0); + auto indices = std::make_shared(node.get_input(1), element::i64); + auto segment_ids = std::make_shared(node.get_input(2), element::i64); + auto data_rank = std::make_shared(std::make_shared(node.get_input(0))); + + // get the last index from segment_ids + auto segments_ids_size = std::make_shared(segment_ids, element::i64); + auto const_one = create_same_type_const(indices, vector{1}, Shape{1}); + auto const_zero = create_same_type_const(indices, vector{0}, Shape{1}); + auto last_idx = std::make_shared(segments_ids_size, const_one); + + // segment_ids are always sorted, so the last index from segment_ids can be used to determine the number of output + // segments. + auto last_segment_idx = std::make_shared(segment_ids, last_idx, const_zero); + auto n_segments = std::make_shared(last_segment_idx, const_one); + + // get sums of sparse segments + auto embedding_segments_sum = + make_shared(data, indices, segment_ids, std::make_shared(n_segments)); + + // get the sizes of each segment + auto unique_segment_ids = make_shared(segment_ids, true, element::i64, element::i64); + auto broadcast = make_shared(const_one, n_segments); + auto divisors = make_shared(broadcast, + unique_segment_ids->output(0), + unique_segment_ids->output(3), + const_zero); + auto divisors_with_correct_type = make_shared(divisors, data); + auto divisors_shape = make_shared(make_shared(const_one, data_rank), + const_zero, + n_segments, + const_zero); + auto divisors_with_correct_shape = std::make_shared(divisors_with_correct_type, divisors_shape, false); + + // divide the sums by the size of the segments + auto mean = std::make_shared(embedding_segments_sum, divisors_with_correct_shape); + + set_node_name(node.get_name(), mean); + return {mean}; +} + +} // namespace op +} // namespace tensorflow +} // namespace frontend +} // namespace ov diff --git a/tests/layer_tests/tensorflow_tests/test_tf_SparseSegmentMean.py b/tests/layer_tests/tensorflow_tests/test_tf_SparseSegmentMean.py new file mode 100644 index 00000000000000..5674a8d88da2b5 --- /dev/null +++ b/tests/layer_tests/tensorflow_tests/test_tf_SparseSegmentMean.py @@ -0,0 +1,79 @@ +# Copyright (C) 2018-2024 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 + +rng = np.random.default_rng(475912) + + +class TestSparseSegmentMean(CommonTFLayerTest): + def _prepare_input(self, inputs_info): + assert 'indices:0' in inputs_info + assert 'values:0' in inputs_info + assert 'segment_indices:0' in inputs_info + + values_shape = inputs_info['values:0'] + + inputs_data = {} + inputs_data['values:0'] = rng.uniform(-5.0, 5.0, values_shape).astype(self.data_type) + + # generate all possible indices + all_indices = [] + for row_ind in range(0, self.shape[0]): + all_indices.append(row_ind) + inputs_data['indices:0'] = rng.choice(all_indices, self.indices_shape[0], replace=False).astype(self.indices_type) + + segment_ids = [] + for ind in range(0, self.indices_shape[0]): + segment_ids.append(self.segment_indices_type(rng.integers(0, self.segments_num))) + inputs_data['segment_indices:0'] = sorted(segment_ids) + + return inputs_data + + def create_sparse_segment_mean(self, data_type, indices_type, segment_indices_type, + shape, indices_shape, segments_num): + self.data_type = data_type + self.indices_type = indices_type + self.segment_indices_type = segment_indices_type + self.shape = shape + self.indices_shape = indices_shape + self.segments_num = segments_num + tf.compat.v1.reset_default_graph() + with tf.compat.v1.Session() as sess: + values = tf.compat.v1.placeholder(data_type, shape, 'values') + indices = tf.compat.v1.placeholder(indices_type, indices_shape, 'indices') + segments_ids = tf.compat.v1.placeholder(segment_indices_type, indices_shape, 'segment_indices') + tf.raw_ops.SparseSegmentMean( + data=values, + indices=indices, + segment_ids=segments_ids) + tf.compat.v1.global_variables_initializer() + tf_net = sess.graph_def + + return tf_net, None + + @pytest.mark.parametrize('data_type', [np.float16, np.float32, np.float64]) + @pytest.mark.parametrize('indices_type', [np.int32, np.int64]) + @pytest.mark.parametrize('segment_indices_type', [np.int32, np.int64]) + @pytest.mark.parametrize('shape, indices_shape, segments_num', [ + [[10], [7], 8], + [[5], [5], 3], + [[5], [2], 4], + [[10, 20], [7], 8], + [[10, 2, 4], [10], 4] + ]) + @pytest.mark.precommit + @pytest.mark.nightly + def test_sparse_segment_mean(self, data_type, indices_type, segment_indices_type, + shape, indices_shape, segments_num, + ie_device, precision, ir_version, temp_dir, + use_legacy_frontend): + if ie_device == 'GPU': + pytest.skip("GPU error: to_shape was called on a dynamic shape, ticket: 152352") + self._test(*self.create_sparse_segment_mean(data_type, indices_type, segment_indices_type, + shape, indices_shape, segments_num), + ie_device, precision, ir_version, temp_dir=temp_dir, + use_legacy_frontend=use_legacy_frontend)