[TF FE] Refactor DepthToSpace and test it (#13544)

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

src/frontends/tensorflow/src/op/depth_to_space.cpp

@@ -15,26 +15,27 @@ namespace tensorflow {
 namespace op {
 
 OutputVector translate_depth_to_space_op(const NodeContext& node) {
-    Output<Node> ng_input = node.get_input(0);
+    default_op_checks(node, 1, {"DepthToSpace"});
+    auto input_data = node.get_input(0);
 
-    // Get the attributes
+    // retrieve attributes
     auto block_size = node.get_attribute<int64_t>("block_size");
-    std::string tf_data_format = node.get_attribute<std::string>("data_format");
+    auto data_format = node.get_attribute<std::string>("data_format");
 
     TENSORFLOW_OP_VALIDATION(node,
-                             tf_data_format == "NHWC" || tf_data_format == "NCHW",
-                             "DepthToSpace data format is neither NHWC nor NCHW");
+                             data_format == "NHWC" || data_format == "NCHW",
+                             "TensorFlow Frontend supports input data for DepthToSpace either in NHWC or NCHW format.");
 
-    bool is_nhwc = (tf_data_format == "NHWC");
+    bool is_nhwc = (data_format == "NHWC");
 
-    convert_nhwc_to_nchw(is_nhwc, ng_input);
-    auto ng_mode = DepthToSpace::DepthToSpaceMode::BLOCKS_FIRST;
-    Output<Node> res = make_shared<DepthToSpace>(ng_input, ng_mode, block_size)->output(0);
-    convert_nchw_to_nhwc(is_nhwc, res);
-    set_node_name(node.get_name(), res.get_node_shared_ptr());
-    return {res};
+    convert_nhwc_to_nchw(is_nhwc, input_data);
+    auto mode = DepthToSpace::DepthToSpaceMode::BLOCKS_FIRST;
+    auto depth_to_space = make_shared<DepthToSpace>(input_data, mode, block_size)->output(0);
+    convert_nchw_to_nhwc(is_nhwc, depth_to_space);
+    set_node_name(node.get_name(), depth_to_space.get_node_shared_ptr());
+    return {depth_to_space};
 }
 }  // namespace op
 }  // namespace tensorflow
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

tests/layer_tests/tensorflow_tests/test_tf_DepthToSpace.py

@@ -0,0 +1,35 @@
+# Copyright (C) 2018-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import tensorflow as tf
+from common.tf_layer_test_class import CommonTFLayerTest
+
+
+class TestDepthToSpace(CommonTFLayerTest):
+    def create_depth_to_space_net(self, input_shape, block_size, data_format):
+        tf.compat.v1.reset_default_graph()
+        # Create the graph and model
+        with tf.compat.v1.Session() as sess:
+            x = tf.compat.v1.placeholder(tf.float32, input_shape, 'x')
+            tf.raw_ops.DepthToSpace(input=x, block_size=block_size, data_format=data_format)
+            tf.compat.v1.global_variables_initializer()
+            tf_net = sess.graph_def
+
+        return tf_net, None
+
+    test_data_basic = [
+        dict(input_shape=[2, 3, 4, 27], block_size=3, data_format='NHWC'),
+        dict(input_shape=[2, 5, 4, 12], block_size=2, data_format='NHWC'),
+        # On CPU TensorFlow does not support 'NCHW' for DepthToSpace
+        # dict(input_shape=[2, 12, 5, 4], block_size=2, data_format='NCHW'),
+    ]
+
+    @pytest.mark.parametrize("params", test_data_basic)
+    @pytest.mark.precommit_tf_fe
+    @pytest.mark.nightly
+    def test_depth_to_space_basic(self, params, ie_device, precision, ir_version, temp_dir,
+                                  use_new_frontend, use_old_api):
+        self._test(*self.create_depth_to_space_net(**params),
+                   ie_device, precision, ir_version, temp_dir=temp_dir,
+                   use_new_frontend=use_new_frontend, use_old_api=use_old_api)