[FusedOps] ShuffleChannels (#2927)

* ShuffleChannels implementation

* Validation of ShuffleChannels params

* Implementation of ShuffleChannels decompose_op()

* Formatting adjustments

* Corrected implementation and validation of op params

* Basic test of ShuffleChannels

* Negative axis value test

* Default params for the ShuffleChannels op

* ShuffleChannels test with floats

* ShuffleChannels validation unit tests

* PR comments

* Compilation error fix

* PR feedback and cleanup

* Code formatting adjustment

* Negative axis value documentation

* Docs update (PR feedback)

* PR feedback: shape and axis validation

* Modify axis semantics on shuffle op

* Revert "PR feedback: shape and axis validation"

This reverts commit 21b708e710b91da2a7e37a69c0da1f31c7743b47.

src/ngraph/CMakeLists.txt

@@ -314,6 +314,8 @@ set (SRC
     op/fused/prelu.hpp
     op/fused/scale_shift.cpp
     op/fused/scale_shift.hpp
+    op/fused/shuffle_channels.cpp
+    op/fused/shuffle_channels.hpp
     op/fused/space_to_depth.cpp
     op/fused/space_to_depth.hpp
     op/fused/split.cpp

src/ngraph/ngraph.hpp

@@ -108,6 +108,7 @@
 #include "ngraph/op/fused/normalize.hpp"
 #include "ngraph/op/fused/prelu.hpp"
 #include "ngraph/op/fused/scale_shift.hpp"
+#include "ngraph/op/fused/shuffle_channels.hpp"
 #include "ngraph/op/fused/space_to_depth.hpp"
 #include "ngraph/op/fused/split.hpp"
 #include "ngraph/op/fused/squared_difference.hpp"

src/ngraph/op/fused/shuffle_channels.cpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+#include "ngraph/op/fused/shuffle_channels.hpp"
+#include "ngraph/op/reshape.hpp"
+#include "ngraph/op/util/reshape.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+op::ShuffleChannels::ShuffleChannels(const shared_ptr<Node>& data,
+                                     const int axis,
+                                     const size_t groups)
+    : FusedOp("ShuffleChannels", {data})
+    , m_axis(axis)
+    , m_groups{groups}
+{
+    constructor_validate_and_infer_types();
+}
+
+size_t op::ShuffleChannels::get_zero_based_axis() const
+{
+    if (m_axis >= 0)
+    {
+        return m_axis;
+    }
+    else
+    {
+        if (!get_input_partial_shape(0).rank().is_dynamic())
+        {
+            return m_axis + static_cast<size_t>(get_input_partial_shape(0).rank());
+        }
+        else
+        {
+            throw ngraph_error("Cannot request zero-based axis with a input of unknown rank");
+        }
+    }
+}
+
+void op::ShuffleChannels::pre_validate_and_infer_types()
+{
+    if (get_input_partial_shape(0).is_static())
+    {
+        const auto shape = get_argument(0)->get_shape();
+
+        NODE_VALIDATION_CHECK(
+            this, shape.size() >= 1, "The input tensor's shape is expected to be at least 1D.");
+        size_t axis_zb = get_zero_based_axis();
+
+        NODE_VALIDATION_CHECK(this,
+                              axis_zb >= 0 && axis_zb < shape.size(),
+                              "The 'axis' parameter for ShuffleChannels has to point to one of the "
+                              "input tensor's shape dimensions.");
+
+        const auto channel_dim_size = shape.at(axis_zb);
+        NODE_VALIDATION_CHECK(
+            this,
+            channel_dim_size % m_groups == 0,
+            "The channel dimension size has to be a multiple of the groups parameter value.");
+    }
+}
+
+NodeVector op::ShuffleChannels::decompose_op() const
+{
+    const auto data = get_argument(0);
+    const auto& data_shape = data->get_shape();
+
+    const auto reshaped = util::reshape(data, get_pre_shuffle_shape(data_shape));
+    const auto shuffled = util::reorder_axes(reshaped, {0, 2, 1, 3});
+
+    return {util::reshape(shuffled, data_shape)};
+}
+
+shared_ptr<Node> op::ShuffleChannels::copy_with_new_args(const NodeVector& new_args) const
+{
+    if (new_args.size() != 1)
+    {
+        throw ngraph_error("Expected 1 element in new_args for the ShuffleChannels op but got " +
+                           std::to_string(new_args.size()));
+    }
+
+    return make_shared<ShuffleChannels>(new_args.at(0), m_axis, m_groups);
+}
+
+Shape op::ShuffleChannels::get_pre_shuffle_shape(const Shape& data_shape) const
+{
+    const Shape& ds = data_shape;
+
+    // in general the resulting shape should contain the following values:
+    // [0]: ds[0] * ds[1] * ... * ds[m_axis-1] (or 1 if m_axis == 0)
+    // [1]: m_groups
+    // [2]: ds[axis] / m_groups
+    // [3]: ds[axis+1] * ds[axis+2] * ... * ds[ds.size()-1] (or 1 if m_axis points to the last elem of ds)
+    Shape res(4, 1);
+
+    size_t axis_zb = get_zero_based_axis();
+    for (size_t i = 0; i < axis_zb; ++i)
+    {
+        res[0] *= ds[i];
+    }
+
+    res[1] = m_groups;
+    res[2] = ds[axis_zb] / m_groups;
+
+    for (size_t i = axis_zb + 1; i < ds.size(); ++i)
+    {
+        res[3] *= ds[i];
+    }
+
+    return res;
+}

src/ngraph/op/fused/shuffle_channels.hpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#pragma once
+
+#include <memory>
+
+#include "ngraph/node.hpp"
+#include "ngraph/op/util/fused_op.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        /// \brief Permutes data in the channel dimension of the input
+        class ShuffleChannels : public ngraph::op::util::FusedOp
+        {
+        public:
+            /// \brief Constructs a ShuffleChannels node.
+            ///
+            /// \param data - Node producing the input tensor
+            /// \param axis - channel dimension index in the data tensor. A negative value means that the index should be calculated from the back of the input data shape.
+            /// \param groups - number of groups the channel dimension specified by axis should be split into
+            ShuffleChannels(const std::shared_ptr<ngraph::Node>& data,
+                            const int axis = 1,
+                            const size_t groups = 1UL);
+
+            size_t get_zero_based_axis() const;
+
+            virtual void pre_validate_and_infer_types() override;
+
+            virtual NodeVector decompose_op() const override;
+
+            virtual std::shared_ptr<Node>
+                copy_with_new_args(const NodeVector& new_args) const override;
+
+            int get_axis() const { return m_axis; }
+            size_t get_groups() const { return m_groups; }
+        private:
+            /// \brief Generates a shape required to permute the data
+            ///
+            /// \param data_shape - Shape of the original input data tensor
+            /// \return A 4D tensor to be used to reshape the input data before shuffling it
+            Shape get_pre_shuffle_shape(const Shape& data_shape) const;
+
+            int m_axis;
+            const size_t m_groups;
+        };
+    }
+}

src/ngraph/op/fused_op_tbl.hpp

@@ -33,6 +33,7 @@ NGRAPH_OP(Normalize, ngraph::op)
 NGRAPH_OP(PRelu, ngraph::op)
 NGRAPH_OP(ScaleShift, ngraph::op)
 NGRAPH_OP(SpaceToDepth, ngraph::op)
+NGRAPH_OP(ShuffleChannels, ngraph::op)
 NGRAPH_OP(SquaredDifference, ngraph::op)
 NGRAPH_OP(Squeeze, ngraph::op)
 NGRAPH_OP(Split, ngraph::op)

src/ngraph/runtime/intelgpu/intelgpu_backend.cpp

@@ -89,6 +89,7 @@
 #include "ngraph/op/fused/mvn.hpp"
 #include "ngraph/op/fused/normalize.hpp"
 #include "ngraph/op/fused/scale_shift.hpp"
+#include "ngraph/op/fused/shuffle_channels.hpp"
 #include "ngraph/op/fused/space_to_depth.hpp"
 #include "ngraph/op/fused/squeeze.hpp"
 #include "ngraph/op/fused/unsqueeze.hpp"
@@ -2079,6 +2080,7 @@ shared_ptr<runtime::Executable>
         case OP_TYPEID::ScatterAdd:
         case OP_TYPEID::ScatterNDAdd:
         case OP_TYPEID::ShapeOf:
+        case OP_TYPEID::ShuffleChannels:
         case OP_TYPEID::SpaceToDepth:
         case OP_TYPEID::Split:
         case OP_TYPEID::SquaredDifference:
@@ -2181,6 +2183,7 @@ bool runtime::intelgpu::IntelGPUBackend::is_supported_impl(const Node& node)
     case OP_TYPEID::Normalize:
     case OP_TYPEID::PRelu:
     case OP_TYPEID::ScaleShift:
+    case OP_TYPEID::ShuffleChannels:
     case OP_TYPEID::SpaceToDepth:
     case OP_TYPEID::Split:
     case OP_TYPEID::SquaredDifference:

src/ngraph/serializer.cpp

@@ -79,6 +79,7 @@
 #include "ngraph/op/fused/normalize.hpp"
 #include "ngraph/op/fused/prelu.hpp"
 #include "ngraph/op/fused/scale_shift.hpp"
+#include "ngraph/op/fused/shuffle_channels.hpp"
 #include "ngraph/op/fused/space_to_depth.hpp"
 #include "ngraph/op/fused/split.hpp"
 #include "ngraph/op/fused/squared_difference.hpp"
@@ -1412,6 +1413,13 @@ static shared_ptr<ngraph::Function>
                 node = make_shared<op::ShapeOf>(args[0]);
                 break;
             }
+            case OP_TYPEID::ShuffleChannels:
+            {
+                const auto axis = node_js.at("axis").get<size_t>();
+                const auto groups = node_js.at("groups").get<size_t>();
+                node = make_shared<op::ShuffleChannels>(args[0], axis, groups);
+                break;
+            }
             case OP_TYPEID::Sigmoid:
             {
                 node = make_shared<op::Sigmoid>(args[0]);
@@ -2199,6 +2207,13 @@ static json write(const Node& n, bool binary_constant_data)
     }
     case OP_TYPEID::ShapeOf: { break;
     }
+    case OP_TYPEID::ShuffleChannels:
+    {
+        const auto tmp = dynamic_cast<const op::ShuffleChannels*>(&n);
+        node["axis"] = tmp->get_axis();
+        node["groups"] = tmp->get_groups();
+        break;
+    }
     case OP_TYPEID::Sigmoid: { break;
     }
     case OP_TYPEID::SigmoidBackprop: { break;

test/backend_fused_op.in.cpp

@@ -867,6 +867,65 @@ NGRAPH_TEST(${BACKEND_NAME}, scale_shift)
     test_case.run();
 }
 
+NGRAPH_TEST(${BACKEND_NAME}, shuffle_channels_simple)
+{
+    const auto data = make_shared<op::Parameter>(element::i32, Shape{1, 15, 2, 2});
+    auto tested_op = make_shared<op::ShuffleChannels>(data, 1, 5);
+    auto function = make_shared<Function>(tested_op, ParameterVector{data});
+
+    auto test_case = ngraph::test::NgraphTestCase(function, "${BACKEND_NAME}");
+
+    std::vector<int32_t> input_data(60);
+    std::iota(std::begin(input_data), std::end(input_data), 0);
+    test_case.add_input(input_data);
+
+    test_case.add_expected_output<int32_t>(
+        Shape{1, 15, 2, 2},
+        {0, 1, 2,  3,  12, 13, 14, 15, 24, 25, 26, 27, 36, 37, 38, 39, 48, 49, 50, 51,
+         4, 5, 6,  7,  16, 17, 18, 19, 28, 29, 30, 31, 40, 41, 42, 43, 52, 53, 54, 55,
+         8, 9, 10, 11, 20, 21, 22, 23, 32, 33, 34, 35, 44, 45, 46, 47, 56, 57, 58, 59});
+
+    test_case.run();
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, shuffle_channels_negative_axis)
+{
+    // in this test the output is the same as in shuffle_channels_simple but
+    // the axis value is negative and the C(channels) value is in a different dimension(0) of the shape
+    const auto data = make_shared<op::Parameter>(element::i32, Shape{15, 2, 1, 2});
+    auto tested_op = make_shared<op::ShuffleChannels>(data, -4, 5);
+    auto function = make_shared<Function>(tested_op, ParameterVector{data});
+
+    auto test_case = ngraph::test::NgraphTestCase(function, "${BACKEND_NAME}");
+
+    std::vector<int32_t> input_data(60);
+    std::iota(std::begin(input_data), std::end(input_data), 0);
+    test_case.add_input(input_data);
+
+    test_case.add_expected_output<int32_t>(
+        Shape{15, 2, 1, 2},
+        {0, 1, 2,  3,  12, 13, 14, 15, 24, 25, 26, 27, 36, 37, 38, 39, 48, 49, 50, 51,
+         4, 5, 6,  7,  16, 17, 18, 19, 28, 29, 30, 31, 40, 41, 42, 43, 52, 53, 54, 55,
+         8, 9, 10, 11, 20, 21, 22, 23, 32, 33, 34, 35, 44, 45, 46, 47, 56, 57, 58, 59});
+
+    test_case.run();
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, shuffle_channels_float)
+{
+    const auto data = make_shared<op::Parameter>(element::f32, Shape{6, 1, 1, 1});
+    auto tested_op = make_shared<op::ShuffleChannels>(data, 0, 2);
+    auto function = make_shared<Function>(tested_op, ParameterVector{data});
+
+    auto test_case = ngraph::test::NgraphTestCase(function, "${BACKEND_NAME}");
+
+    test_case.add_input<float>({0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f});
+
+    test_case.add_expected_output<float>(Shape{6, 1, 1, 1}, {0.0f, 3.0f, 1.0f, 4.0f, 2.0f, 5.0f});
+
+    test_case.run();
+}
+
 NGRAPH_TEST(${BACKEND_NAME}, squeeze)
 {
     const auto data_node = make_shared<op::Parameter>(element::f32, Shape{1, 4, 1, 1, 2});

test/type_prop.cpp

@@ -14568,6 +14568,65 @@ TEST(type_prop, scale_shift)
     EXPECT_EQ(scale_shift_func->get_shape(), (Shape{3, 6}));
 }
 
+TEST(type_prop, shuffle_channels_axis_validation)
+{
+    try
+    {
+        const auto data = make_shared<op::Parameter>(element::f64, Shape{1, 2, 3, 4});
+        const auto shuffle_channels = make_shared<op::ShuffleChannels>(data, -5, 5);
+        FAIL() << "ShuffleChannels validation did not work. Op node was created with incorrect "
+                  "params.";
+    }
+    catch (const NodeValidationFailure& error)
+    {
+        EXPECT_HAS_SUBSTRING(error.what(),
+                             "The 'axis' parameter for ShuffleChannels has to point to one of the "
+                             "input tensor's shape dimensions");
+    }
+}
+
+TEST(type_prop, shuffle_channels_negative_axis_calculation)
+{
+    const auto data = make_shared<op::Parameter>(element::f64, Shape{1, 2, 3, 4});
+
+    const auto shuffle_channels = make_shared<op::ShuffleChannels>(data, -3, 2);
+
+    EXPECT_EQ(shuffle_channels->get_zero_based_axis(), 1);
+}
+
+TEST(type_prop, shuffle_channels_invalid_input_shape)
+{
+    try
+    {
+        const auto data = make_shared<op::Parameter>(element::f64, Shape{});
+        const auto shuffle_channels = make_shared<op::ShuffleChannels>(data, 0, 1);
+        FAIL() << "ShuffleChannels validation did not work. Op node was created with incorrect "
+                  "params.";
+    }
+    catch (const NodeValidationFailure& error)
+    {
+        EXPECT_HAS_SUBSTRING(error.what(),
+                             "The input tensor's shape is expected to be at least 1D.");
+    }
+}
+
+TEST(type_prop, shuffle_channels_invalid_groups_value)
+{
+    try
+    {
+        const auto data = make_shared<op::Parameter>(element::f64, Shape{1, 2, 3, 15});
+        const auto shuffle_channels = make_shared<op::ShuffleChannels>(data, -1, 2);
+        FAIL() << "ShuffleChannels validation did not work. Op node was created with incorrect "
+                  "params.";
+    }
+    catch (const NodeValidationFailure& error)
+    {
+        EXPECT_HAS_SUBSTRING(
+            error.what(),
+            "The channel dimension size has to be a multiple of the groups parameter value.");
+    }
+}
+
 TEST(type_prop, squared_difference)
 {
     const auto x1 = make_shared<op::Parameter>(element::f64, Shape{2, 2});