From 089fc0d5c1ee77fb13457a2e7140ccc357acfcda Mon Sep 17 00:00:00 2001 From: Roman Kazantsev Date: Mon, 26 Feb 2024 23:45:24 +0400 Subject: [PATCH] [TF FE] MatrixBandPart operation for TensorFlow Hub models (#23082) **Details:** `MatrixBandPart` is needed to support Keras StableDiffusion model. This is reserved PR for https://github.com/openvinotoolkit/openvino/pull/22447 **Ticket:** CVS-133786 --------- Signed-off-by: Kazantsev, Roman Co-authored-by: himanshugupta11002 <72141497+himanshugupta11002@users.noreply.github.com> --- .../tensorflow/docs/supported_ops.md | 2 +- src/frontends/tensorflow/src/op_table.cpp | 1 + .../include/common_op_table.hpp | 1 + .../src/op/matrix_band_part.cpp | 90 +++++++++++++++++++ .../test_tf_MatrixBandPart.py | 41 +++++++++ 5 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 src/frontends/tensorflow_common/src/op/matrix_band_part.cpp create mode 100644 tests/layer_tests/tensorflow_tests/test_tf_MatrixBandPart.py diff --git a/src/frontends/tensorflow/docs/supported_ops.md b/src/frontends/tensorflow/docs/supported_ops.md index 94956372a0bd94..3e8fa158c6c4a6 100644 --- a/src/frontends/tensorflow/docs/supported_ops.md +++ b/src/frontends/tensorflow/docs/supported_ops.md @@ -620,7 +620,7 @@ A "supported operation" is one that TensorFlow Frontend can convert to the OpenV | MatMul | YES | | | MatchingFiles | NO | | | MatchingFilesDataset | NO | | -| MatrixBandPart | NO | | +| MatrixBandPart | YES | | | MatrixDeterminant | NO | | | MatrixDiag | YES | | | MatrixDiagPart | NO | | diff --git a/src/frontends/tensorflow/src/op_table.cpp b/src/frontends/tensorflow/src/op_table.cpp index 034861ba6030a5..bdc13df7b0e08a 100644 --- a/src/frontends/tensorflow/src/op_table.cpp +++ b/src/frontends/tensorflow/src/op_table.cpp @@ -298,6 +298,7 @@ const std::map get_supported_ops() { {"LookupTableInsertV2", CreatorFunction(translate_no_op)}, {"LRN", CreatorFunction(translate_lrn_op)}, {"MatMul", CreatorFunction(translate_mat_mul_op)}, + {"MatrixBandPart", CreatorFunction(translate_matrix_band_part_op)}, {"MatrixDiag", CreatorFunction(translate_matrix_diag_op)}, {"MaxPool", CreatorFunction(translate_max_pool_op)}, {"MaxPoolV2", CreatorFunction(translate_max_pool_op)}, diff --git a/src/frontends/tensorflow_common/include/common_op_table.hpp b/src/frontends/tensorflow_common/include/common_op_table.hpp index 9a4376951c14cf..f0b564f0f07a86 100644 --- a/src/frontends/tensorflow_common/include/common_op_table.hpp +++ b/src/frontends/tensorflow_common/include/common_op_table.hpp @@ -102,6 +102,7 @@ OP_CONVERTER(translate_log_1p_op); OP_CONVERTER(translate_lrn_op); OP_CONVERTER(translate_mat_mul_op); OP_CONVERTER(translate_matrix_diag_op); +OP_CONVERTER(translate_matrix_band_part_op); OP_CONVERTER(translate_max_pool_op); OP_CONVERTER_NAMED(translate_max_pool_with_argmax); OP_CONVERTER(translate_mirror_pad_op); diff --git a/src/frontends/tensorflow_common/src/op/matrix_band_part.cpp b/src/frontends/tensorflow_common/src/op/matrix_band_part.cpp new file mode 100644 index 00000000000000..f542f7fa256795 --- /dev/null +++ b/src/frontends/tensorflow_common/src/op/matrix_band_part.cpp @@ -0,0 +1,90 @@ +// Copyright (C) 2024 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 + +#include "openvino/op/constant.hpp" +#include "openvino/op/convert.hpp" +#include "openvino/op/less.hpp" +#include "openvino/op/less_eq.hpp" +#include "openvino/op/logical_and.hpp" +#include "openvino/op/logical_or.hpp" +#include "openvino/op/range.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/select.hpp" +#include "openvino/op/shape_of.hpp" +#include "openvino/op/slice.hpp" +#include "openvino/op/subtract.hpp" +#include "openvino/op/unsqueeze.hpp" +#include "utils.hpp" + +using namespace std; +using namespace ov::op; + +namespace ov { +namespace frontend { +namespace tensorflow { +namespace op { + +OutputVector translate_matrix_band_part_op(const NodeContext& node) { + default_op_checks(node, 3, {"MatrixBandPart"}); + + // Input tensor and parameters + auto input = node.get_input(0); + auto num_lower = node.get_input(1); + auto num_upper = node.get_input(2); + + // create scalar auxiliary constants + auto const_zero = make_shared(element::i64, Shape{}, 0); + auto const_one = make_shared(element::i64, Shape{}, 1); + auto const_two = make_shared(element::i64, Shape{}, 2); + + // input has a shape [I, J, K, ..., M, N] + // compute sizes of two last dimensions of M and N + auto input_shape = make_shared(input, element::i64); + auto input_rank = make_shared(input_shape, element::i64); + auto input_rank_minus_one = make_shared(input_rank, const_one); + auto input_rank_minus_two = make_shared(input_rank, const_two); + auto slice_step = make_shared(element::i64, Shape{1}, 1); + auto slice_axis = make_shared(element::i64, Shape{1}, 0); + auto m = make_shared(input_shape, input_rank_minus_two, input_rank_minus_one, slice_step, slice_axis) + ->output(0); + auto n = make_shared(input_shape, input_rank_minus_one, input_rank, slice_step, slice_axis)->output(0); + + // generate ranges [0, M) and [0, N) + auto scalar_shape = make_shared(element::i64, Shape{0}, vector{}); + m = make_shared(m, scalar_shape, false); + n = make_shared(n, scalar_shape, false); + auto range_m = make_shared(const_zero, m, const_one, element::i64)->output(0); + auto range_n = make_shared(const_zero, n, const_one, element::i64)->output(0); + range_m = make_shared(range_m, const_one); + range_n = make_shared(range_n, const_zero); + + // adjust num_lower and num_upper to have them of type i64 + // the same as M and N + // it is needed for in_band computation + num_lower = make_shared(num_lower, element::i64); + num_upper = make_shared(num_upper, element::i64); + + // compute in_band(m, n) = (num_lower < 0 || (m-n) <= num_lower)) && (num_upper < 0 || (n-m) <= num_upper) + auto num_lower_less_zero = make_shared(num_lower, const_zero); + auto i_minus_j = make_shared(range_m, range_n); + auto i_minus_j_less_eq_num_lower = make_shared(i_minus_j, num_lower); + auto num_upper_less_zero = make_shared(num_upper, const_zero); + auto j_minus_i = make_shared(range_n, range_m); + auto j_minus_i_less_eq_num_upper = make_shared(j_minus_i, num_upper); + auto in_band1 = make_shared(num_lower_less_zero, i_minus_j_less_eq_num_lower); + auto in_band2 = make_shared(num_upper_less_zero, j_minus_i_less_eq_num_upper); + auto in_band = make_shared(in_band1, in_band2); + + // create zero constant of the same type as input + auto zero = create_same_type_const_scalar(input, 0); + + auto result = make_shared(in_band, input, zero); + + set_node_name(node.get_name(), result); + return {result}; +} + +} // namespace op +} // namespace tensorflow +} // namespace frontend +} // namespace ov diff --git a/tests/layer_tests/tensorflow_tests/test_tf_MatrixBandPart.py b/tests/layer_tests/tensorflow_tests/test_tf_MatrixBandPart.py new file mode 100644 index 00000000000000..10f573cef604a3 --- /dev/null +++ b/tests/layer_tests/tensorflow_tests/test_tf_MatrixBandPart.py @@ -0,0 +1,41 @@ +# Copyright (C) 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() + + +class TestMatrixBandPart(CommonTFLayerTest): + def _prepare_input(self, inputs_info): + assert 'input:0' in inputs_info + input_shape = inputs_info['input:0'] + inputs_data = {} + inputs_data['input:0'] = rng.integers(-8, 8, input_shape).astype(self.input_type) + return inputs_data + + def create_matrix_band_part_net(self, input_shape, input_type, num_lower, num_upper): + self.input_type = input_type + tf.compat.v1.reset_default_graph() + with tf.compat.v1.Session() as sess: + input_tensor = tf.compat.v1.placeholder(input_type, input_shape, 'input') + tf.raw_ops.MatrixBandPart(input=input_tensor, num_lower=num_lower, num_upper=num_upper) + tf.compat.v1.global_variables_initializer() + tf_net = sess.graph_def + return tf_net, None + + @pytest.mark.parametrize('input_shape', [[5, 5], [3, 4, 4], [1, 2, 5, 5], [3, 5, 4]]) + @pytest.mark.parametrize('input_type', [np.float32, np.int32]) + @pytest.mark.parametrize('num_lower', [-4, -1, 0, 1, 4]) + @pytest.mark.parametrize('num_upper', [-4, -1, 0, 1, 4]) + @pytest.mark.precommit_tf_fe + @pytest.mark.nightly + def test_matrix_band_part_basic(self, input_shape, input_type, num_lower, num_upper, + ie_device, precision, ir_version, temp_dir, + use_legacy_frontend): + self._test(*self.create_matrix_band_part_net(input_shape, input_type, num_lower, num_upper), + ie_device, precision, ir_version, temp_dir=temp_dir, + use_legacy_frontend=use_legacy_frontend)