[Fused] Squeeze op with dynamic input (#2896)

* Add fused Squeeze op

* Use fused Squeeze op in ONNX importer

* Update serializer

* Add unit tests

* Add type prop tests

* Change Squeeze signature to accept a dynamic input for axes

* Update src/ngraph/op/fused/squeeze.cpp
Co-Authored-By: Adam Rogowiec <adam.rogowiec@intel.com>

* Code review comments

* Fix failing unit test

* Code review comments

* Code review comments

* style

* Add op to iGPU backend

* Add op to iGPU backend

src/ngraph/CMakeLists.txt

@@ -308,6 +308,8 @@ set (SRC
     op/fused/scale_shift.hpp
     op/fused/space_to_depth.cpp
     op/fused/space_to_depth.hpp
+    op/fused/squeeze.cpp
+    op/fused/squeeze.hpp
     op/util/arithmetic_reduction.cpp
     op/util/arithmetic_reduction.hpp
     op/util/binary_elementwise_arithmetic.cpp

src/ngraph/frontend/onnx_import/op/squeeze.cpp

@@ -14,21 +14,12 @@
 // limitations under the License.
 //*****************************************************************************
 
-#include <cstddef>
-#include <functional>
-#include <iterator>
-#include <numeric>
-#include <set>
 #include <vector>
 
 #include "exceptions.hpp"
-#include "ngraph/axis_vector.hpp"
-#include "ngraph/op/reshape.hpp"
-#include "ngraph/op/util/reshape.hpp"
-#include "ngraph/shape.hpp"
-#include "ngraph/util.hpp"
+#include "ngraph/op/constant.hpp"
+#include "ngraph/op/fused/squeeze.hpp"
 #include "squeeze.hpp"
-#include "utils/reshape.hpp"
 
 namespace ngraph
 {
@@ -40,49 +31,19 @@ namespace ngraph
             {
                 NodeVector squeeze(const Node& node)
                 {
-                    NodeVector inputs{node.get_ng_inputs()};
-                    auto data = inputs.at(0);
-                    auto data_shape = data->get_shape();
-                    auto axes = node.get_attribute_value<std::vector<std::size_t>>("axes", {});
-                    AxisVector input_order{ngraph::get_default_order(data_shape.size())};
+                    auto data = node.get_ng_inputs().at(0);
+                    auto axes = node.get_attribute_value<std::vector<std::int64_t>>("axes", {});
 
-                    // Prepare set of unique axes marked to be removed from input data.
-                    if (axes.empty())
+                    for (auto axis : axes)
                     {
-                        // Default behaviour is to remove all single dimension axes.
-                        for (std::size_t idx = 0; idx < data_shape.size(); ++idx)
-                        {
-                            if (data_shape.at(idx) == 1)
-                            {
-                                // Mark with zero elements to remove;
-                                data_shape.at(idx) = 0;
-                            }
-                        }
-                    }
-                    else
-                    {
-                        std::set<std::size_t, std::greater<std::size_t>> unique_axes(
-                            std::begin(axes), std::end(axes));
-                        for (uint64_t axis : unique_axes)
-                        {
-                            ASSERT_VALID_ARGUMENT(node, data_shape.at(axis) == 1)
-                                << "provided axis value is invalid. Only single dimension axes may "
-                                   "be removed.";
-                            // Mark with zero elements to remove;
-                            data_shape.at(axis) = 0;
-                        }
+                        ASSERT_VALID_ARGUMENT(node, axis >= 0)
+                            << "provided axes attribute is invalid. Only non-negative "
+                            << "integers are allowed, got " << axis << ".";
                     }
 
-                    Shape output_data_shape;
-                    for (std::size_t idx = 0; idx < data_shape.size(); ++idx)
-                    {
-                        if (data_shape.at(idx) != 0)
-                        {
-                            output_data_shape.push_back(data_shape.at(idx));
-                        }
-                    }
-                    return {std::make_shared<ngraph::op::Reshape>(
-                        data, input_order, output_data_shape)};
+                    auto axes_node = std::make_shared<ngraph::op::Constant>(
+                        element::u64, Shape{axes.size()}, axes);
+                    return {std::make_shared<ngraph::op::Squeeze>(data, axes_node)};
                 }
 
             } // namespace set_1

src/ngraph/ngraph.hpp

@@ -107,6 +107,7 @@
 #include "ngraph/op/fused/prelu.hpp"
 #include "ngraph/op/fused/scale_shift.hpp"
 #include "ngraph/op/fused/space_to_depth.hpp"
+#include "ngraph/op/fused/squeeze.hpp"
 #include "ngraph/op/gather.hpp"
 #include "ngraph/op/gather_nd.hpp"
 #include "ngraph/op/get_output_element.hpp"

src/ngraph/op/fused/squeeze.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 <cstddef>
+#include <functional>
+#include <iterator>
+#include <set>
+
+#include "ngraph/builder/make_constant.hpp"
+#include "ngraph/op/constant.hpp"
+#include "ngraph/op/fused/squeeze.hpp"
+#include "ngraph/op/reshape.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+op::Squeeze::Squeeze(const shared_ptr<Node>& data, const shared_ptr<Node>& axes)
+    : FusedOp("Squeeze", {data, axes})
+{
+    constructor_validate_and_infer_types();
+}
+
+NodeVector op::Squeeze::decompose_op() const
+{
+    auto data = get_argument(0);
+    auto axes_node = get_argument(1);
+
+    // Currently only support Constant node for axes.
+    NODE_VALIDATION_CHECK(this,
+                          axes_node->is_constant(),
+                          "doesn't support 'axes' input of other type than a Constant.");
+
+    // Get value of axes from Constant
+    auto axes_constant = dynamic_pointer_cast<op::Constant>(axes_node);
+    auto axes = axes_constant->get_vector<size_t>();
+
+    auto data_shape = data->get_shape();
+
+    // Prepare set of unique axes marked to be removed from input data.
+    if (axes.empty())
+    {
+        // Default behaviour is to remove all single dimension axes.
+        for (size_t idx = 0; idx < data_shape.size(); ++idx)
+        {
+            if (data_shape.at(idx) == 1)
+            {
+                // Mark with zero elements to remove;
+                data_shape.at(idx) = 0;
+            }
+        }
+    }
+    else
+    {
+        set<size_t, greater<size_t>> unique_axes(begin(axes), end(axes));
+        for (uint64_t axis : unique_axes)
+        {
+            NODE_VALIDATION_CHECK(
+                this,
+                (data_shape.at(axis) == 1),
+                "provided axis value is invalid. Only axes of size 1 may be removed.");
+
+            // Mark with zero elements to remove;
+            data_shape.at(axis) = 0;
+        }
+    }
+
+    Shape output_data_shape;
+    for (size_t idx = 0; idx < data_shape.size(); ++idx)
+    {
+        if (data_shape.at(idx) != 0)
+        {
+            output_data_shape.push_back(data_shape.at(idx));
+        }
+    }
+
+    AxisVector input_order{get_default_order(data_shape.size())};
+    return {make_shared<op::Reshape>(data, input_order, output_data_shape)};
+}
+
+shared_ptr<Node> op::Squeeze::copy_with_new_args(const NodeVector& new_args) const
+{
+    if (new_args.size() != 2)
+    {
+        throw ngraph_error("Incorrect number of new arguments");
+    }
+    return make_shared<Squeeze>(new_args.at(0), new_args.at(1));
+}

src/ngraph/op/fused/squeeze.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/axis_vector.hpp"
+#include "ngraph/node.hpp"
+#include "ngraph/op/op.hpp"
+#include "ngraph/op/util/fused_op.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        class Squeeze : public ngraph::op::util::FusedOp
+        {
+        public:
+            Squeeze(const std::shared_ptr<ngraph::Node>& data,
+                    const std::shared_ptr<ngraph::Node>& axes);
+
+            virtual NodeVector decompose_op() const override;
+
+            virtual std::shared_ptr<Node>
+                copy_with_new_args(const NodeVector& new_args) const override;
+        };
+    }
+}

src/ngraph/op/fused_op_tbl.hpp

@@ -21,13 +21,14 @@ NGRAPH_OP(Clamp, ngraph::op)
 NGRAPH_OP(ConvolutionBias, ngraph::op)
 NGRAPH_OP(ConvolutionBiasAdd, ngraph::op)
 NGRAPH_OP(ConvolutionBiasBackpropFiltersBias, ngraph::op)
-NGRAPH_OP(HardSigmoid, ngraph::op)
 NGRAPH_OP(DepthToSpace, ngraph::op)
 NGRAPH_OP(Elu, ngraph::op)
 NGRAPH_OP(Gemm, ngraph::op)
 NGRAPH_OP(GroupConvolution, ngraph::op)
+NGRAPH_OP(HardSigmoid, ngraph::op)
 NGRAPH_OP(MVN, ngraph::op)
 NGRAPH_OP(Normalize, ngraph::op)
 NGRAPH_OP(PRelu, ngraph::op)
 NGRAPH_OP(ScaleShift, ngraph::op)
 NGRAPH_OP(SpaceToDepth, ngraph::op)
+NGRAPH_OP(Squeeze, ngraph::op)

src/ngraph/runtime/intelgpu/intelgpu_backend.cpp

@@ -88,6 +88,7 @@
 #include "ngraph/op/fused/normalize.hpp"
 #include "ngraph/op/fused/scale_shift.hpp"
 #include "ngraph/op/fused/space_to_depth.hpp"
+#include "ngraph/op/fused/squeeze.hpp"
 #include "ngraph/op/get_output_element.hpp"
 #include "ngraph/op/greater.hpp"
 #include "ngraph/op/greater_eq.hpp"
@@ -2072,6 +2073,7 @@ shared_ptr<runtime::Executable>
         case OP_TYPEID::ScatterNDAdd:
         case OP_TYPEID::ShapeOf:
         case OP_TYPEID::SpaceToDepth:
+        case OP_TYPEID::Squeeze:
         case OP_TYPEID::StopGradient:
         case OP_TYPEID::Tile:
         case OP_TYPEID::Transpose:
@@ -2164,7 +2166,8 @@ bool runtime::intelgpu::IntelGPUBackend::is_supported_impl(const Node& node)
     case OP_TYPEID::MVN:
     case OP_TYPEID::Normalize:
     case OP_TYPEID::PRelu:
-    case OP_TYPEID::SpaceToDepth: { return false;
+    case OP_TYPEID::SpaceToDepth:
+    case OP_TYPEID::Squeeze: { return false;
     }
     default: { return true;
     }

src/ngraph/serializer.cpp

@@ -78,6 +78,7 @@
 #include "ngraph/op/fused/prelu.hpp"
 #include "ngraph/op/fused/scale_shift.hpp"
 #include "ngraph/op/fused/space_to_depth.hpp"
+#include "ngraph/op/fused/squeeze.hpp"
 #include "ngraph/op/gather.hpp"
 #include "ngraph/op/gather_nd.hpp"
 #include "ngraph/op/get_output_element.hpp"
@@ -1443,6 +1444,11 @@ static shared_ptr<ngraph::Function>
                 node = make_shared<op::Sqrt>(args[0]);
                 break;
             }
+            case OP_TYPEID::Squeeze:
+            {
+                node = make_shared<op::Squeeze>(args[0], args[1]);
+                break;
+            }
             case OP_TYPEID::Subtract:
             {
                 node = make_shared<op::Subtract>(args[0], args[1]);
@@ -2173,6 +2179,8 @@ static json write(const Node& n, bool binary_constant_data)
     }
     case OP_TYPEID::Sqrt: { break;
     }
+    case OP_TYPEID::Squeeze: { break;
+    }
     case OP_TYPEID::StopGradient: { break;
     }
     case OP_TYPEID::Subtract: { break;

test/backend_fused_op.in.cpp

@@ -24,6 +24,7 @@
 #include <string>
 
 #include "gtest/gtest.h"
+#include "ngraph/check.hpp"
 #include "ngraph/ngraph.hpp"
 #include "util/all_close.hpp"
 #include "util/all_close_f.hpp"
@@ -757,3 +758,42 @@ NGRAPH_TEST(${BACKEND_NAME}, scale_shift)
     test_case.add_expected_output<double>(Shape{3, 6}, vector<double>(18, 6));
     test_case.run();
 }
+
+NGRAPH_TEST(${BACKEND_NAME}, squeeze)
+{
+    const auto data_node = make_shared<op::Parameter>(element::f32, Shape{1, 4, 1, 1, 2});
+    const auto axes_node =
+        make_shared<ngraph::op::Constant>(element::u64, Shape{2}, vector<int64_t>{0, 2});
+    const auto squeeze = make_shared<op::Squeeze>(data_node, axes_node);
+
+    const auto function = make_shared<Function>(NodeVector{squeeze}, ParameterVector{data_node});
+    auto test_case = ngraph::test::NgraphTestCase(function, "${BACKEND_NAME}");
+
+    const auto data = vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
+    test_case.add_input(data);
+    test_case.add_expected_output<float>(Shape{4, 1, 2}, data);
+    test_case.run();
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, squeeze_default_axes)
+{
+    const auto data_node = make_shared<op::Parameter>(element::f32, Shape{1, 4, 1, 1, 2});
+    const auto axes_node =
+        make_shared<ngraph::op::Constant>(element::u64, Shape{0}, vector<int64_t>{});
+    const auto squeeze = make_shared<op::Squeeze>(data_node, axes_node);
+
+    const auto function = make_shared<Function>(NodeVector{squeeze}, ParameterVector{data_node});
+    auto test_case = ngraph::test::NgraphTestCase(function, "${BACKEND_NAME}");
+
+    const auto data = vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
+    test_case.add_input(data);
+    test_case.add_expected_output<float>(Shape{4, 2}, data);
+    test_case.run();
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, squeeze_dynamic)
+{
+    const auto data_param = make_shared<op::Parameter>(element::f32, Shape{1, 4, 1, 1, 2});
+    const auto axes_param = make_shared<op::Parameter>(element::i64, Shape{2});
+    EXPECT_THROW(make_shared<op::Squeeze>(data_param, axes_param), CheckFailure);
+}

test/type_prop.cpp

@@ -13554,6 +13554,23 @@ TEST(type_prop, space_to_depth)
     ASSERT_EQ(space_to_depth->get_shape(), (Shape{1, 128, 8, 8}));
 }
 
+TEST(type_prop, squeeze)
+{
+    auto param = make_shared<op::Parameter>(element::f32, Shape{1, 4, 1, 4, 1, 8});
+    auto axes_node =
+        make_shared<ngraph::op::Constant>(element::u64, Shape{2}, vector<int64_t>{0, 2});
+    auto squeeze = make_shared<op::Squeeze>(param, axes_node);
+
+    ASSERT_EQ(squeeze->get_element_type(), element::f32);
+    ASSERT_EQ(squeeze->get_shape(), (Shape{4, 4, 1, 8}));
+
+    axes_node = make_shared<ngraph::op::Constant>(element::u64, Shape{0}, vector<int64_t>{});
+    auto squeeze_default_axes = make_shared<op::Squeeze>(param, axes_node);
+
+    ASSERT_EQ(squeeze_default_axes->get_element_type(), element::f32);
+    ASSERT_EQ(squeeze_default_axes->get_shape(), (Shape{4, 4, 8}));
+}
+
 TEST(type_prop, gather_nd_scalar_from_2d)
 {
     Shape params_shape{2, 2};