Adding ROI pooling layers (#2899)

src/ngraph/CMakeLists.txt

@@ -176,10 +176,14 @@ set (SRC
     op/experimental/layers/prior_box_clustered.hpp
     op/experimental/layers/proposal.hpp
     op/experimental/layers/proposal.cpp
+    op/experimental/layers/psroi_pooling.hpp
+    op/experimental/layers/psroi_pooling.cpp
     op/experimental/layers/region_yolo.hpp
     op/experimental/layers/region_yolo.cpp
     op/experimental/layers/reorg_yolo.hpp
     op/experimental/layers/reorg_yolo.cpp
+    op/experimental/layers/roi_pooling.hpp
+    op/experimental/layers/roi_pooling.cpp
     op/floor.cpp
     op/floor.hpp
     op/gather.cpp

src/ngraph/op/experimental/layers/psroi_pooling.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 "psroi_pooling.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+op::PSROIPooling::PSROIPooling(const shared_ptr<Node>& input,
+                               const std::shared_ptr<Node>& coords,
+                               const size_t output_dim,
+                               const size_t group_size,
+                               const float spatial_scale,
+                               const Shape& num_bins,
+                               const std::string& kind)
+    : Op("PSROIPooling", check_single_output_args({input, coords}))
+    , m_output_dim(output_dim)
+    , m_group_size(group_size)
+    , m_spatial_scale(spatial_scale)
+    , m_num_bins(num_bins)
+    , m_kind(kind)
+{
+    constructor_validate_and_infer_types();
+}
+
+void op::PSROIPooling::validate_and_infer_types()
+{
+    auto input_et = get_input_element_type(0);
+    if (get_input_partial_shape(0).is_static() && get_input_partial_shape(1).is_static())
+    {
+        Shape input_shape = get_input_partial_shape(0).to_shape();
+        Shape coords_shape = get_input_partial_shape(1).to_shape();
+        NODE_VALIDATION_CHECK(this,
+                              input_shape.size() >= 3,
+                              "PSROIPooling expects 3 or higher dimensions for input. Got ",
+                              input_shape.size());
+        NODE_VALIDATION_CHECK(this,
+                              coords_shape.size() == 2,
+                              "PSROIPooling expects 2 dimensions for box coordinates. Got ",
+                              coords_shape.size());
+        Shape output_shape{coords_shape[0], m_output_dim};
+        for (size_t i = 2; i < input_shape.size(); i++)
+        {
+            output_shape.push_back(m_group_size);
+        }
+        set_output_type(0, input_et, output_shape);
+    }
+    else
+    {
+        set_output_type(0, input_et, PartialShape::dynamic());
+    }
+}
+
+shared_ptr<Node> op::PSROIPooling::copy_with_new_args(const NodeVector& new_args) const
+{
+    check_new_args_count(this, new_args);
+    return make_shared<PSROIPooling>(new_args.at(0),
+                                     new_args.at(1),
+                                     m_output_dim,
+                                     m_group_size,
+                                     m_spatial_scale,
+                                     m_num_bins,
+                                     m_kind);
+}

src/ngraph/op/experimental/layers/psroi_pooling.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 "ngraph/op/op.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        class PSROIPooling : public Op
+        {
+        public:
+            /// \brief Constructs a PSROIPooling operation
+            ///
+            /// \param input          Input feature map {N, C, ...}
+            /// \param coords         Coordinates of bounding boxes
+            /// \param output_dim     Output channel number
+            /// \param group_size     Number of groups to encode position-sensitive scores
+            /// \param spatial_scale  Ratio of input feature map over input image size
+            /// \param num_bins       Number of bins to divide the input feature maps
+            /// \param kind           Kind of pooling - Avg or Bilinear
+            PSROIPooling(const std::shared_ptr<Node>& input,
+                         const std::shared_ptr<Node>& coords,
+                         const size_t output_dim,
+                         const size_t group_size,
+                         const float spatial_scale,
+                         const Shape& num_bins,
+                         const std::string& kind);
+
+            void validate_and_infer_types() override;
+
+            virtual std::shared_ptr<Node>
+                copy_with_new_args(const NodeVector& new_args) const override;
+
+        private:
+            size_t m_output_dim;
+            size_t m_group_size;
+            float m_spatial_scale;
+            Shape m_num_bins;
+            std::string m_kind;
+        };
+    }
+}

src/ngraph/op/experimental/layers/roi_pooling.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 "roi_pooling.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+op::ROIPooling::ROIPooling(const shared_ptr<Node>& input,
+                           const std::shared_ptr<Node>& coords,
+                           const Shape& output_size,
+                           const float spatial_scale,
+                           const std::string& kind)
+    : Op("ROIPooling", check_single_output_args({input, coords}))
+    , m_output_size(output_size)
+    , m_spatial_scale(spatial_scale)
+    , m_kind(kind)
+{
+    constructor_validate_and_infer_types();
+}
+
+void op::ROIPooling::validate_and_infer_types()
+{
+    auto input_et = get_input_element_type(0);
+    if (get_input_partial_shape(0).is_static() && get_input_partial_shape(1).is_static())
+    {
+        Shape input_shape = get_input_partial_shape(0).to_shape();
+        Shape coords_shape = get_input_partial_shape(1).to_shape();
+        NODE_VALIDATION_CHECK(this,
+                              input_shape.size() >= 3,
+                              "ROIPooling expects 3 or higher dimensions for input. Got ",
+                              input_shape.size());
+        NODE_VALIDATION_CHECK(this,
+                              coords_shape.size() == 2,
+                              "ROIPooling expects 2 dimensions for box coordinates. Got ",
+                              coords_shape.size());
+        NODE_VALIDATION_CHECK(this,
+                              input_shape.size() - 2 == m_output_size.size(),
+                              "Spatial dimensions on input: ",
+                              input_shape.size() - 2,
+                              " doesn't match dimensions on requested output_size: ",
+                              m_output_size.size());
+        Shape output_shape{coords_shape[0], input_shape[1]};
+        output_shape.insert(output_shape.end(), m_output_size.begin(), m_output_size.end());
+        set_output_type(0, input_et, output_shape);
+    }
+    else
+    {
+        set_output_type(0, input_et, PartialShape::dynamic());
+    }
+}
+
+shared_ptr<Node> op::ROIPooling::copy_with_new_args(const NodeVector& new_args) const
+{
+    check_new_args_count(this, new_args);
+    return make_shared<ROIPooling>(
+        new_args.at(0), new_args.at(1), m_output_size, m_spatial_scale, m_kind);
+}

src/ngraph/op/experimental/layers/roi_pooling.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 "ngraph/op/op.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        class ROIPooling : public Op
+        {
+        public:
+            /// \brief Constructs a ROIPooling operation
+            ///
+            /// \param input          Input feature map {N, C, ...}
+            /// \param coords         Coordinates of bounding boxes
+            /// \param output_size    Height/Width of ROI output features
+            /// \param spatial_scale  Ratio of input feature map over input image size
+            /// \param kind           Kind of pooling - Max or Bilinear
+            ROIPooling(const std::shared_ptr<Node>& input,
+                       const std::shared_ptr<Node>& coords,
+                       const Shape& output_size,
+                       const float spatial_scale,
+                       const std::string& kind);
+
+            void validate_and_infer_types() override;
+
+            virtual std::shared_ptr<Node>
+                copy_with_new_args(const NodeVector& new_args) const override;
+
+        private:
+            Shape m_output_size;
+            float m_spatial_scale;
+            std::string m_kind;
+        };
+    }
+}

test/type_prop_layers.cpp

@@ -22,8 +22,10 @@
 #include "ngraph/op/experimental/layers/prior_box.hpp"
 #include "ngraph/op/experimental/layers/prior_box_clustered.hpp"
 #include "ngraph/op/experimental/layers/proposal.hpp"
+#include "ngraph/op/experimental/layers/psroi_pooling.hpp"
 #include "ngraph/op/experimental/layers/region_yolo.hpp"
 #include "ngraph/op/experimental/layers/reorg_yolo.hpp"
+#include "ngraph/op/experimental/layers/roi_pooling.hpp"
 
 #include <memory>
 using namespace std;
@@ -203,3 +205,19 @@ TEST(type_prop_layers, reorg_yolo)
     auto op = make_shared<op::ReorgYolo>(inputs, Strides{2});
     ASSERT_EQ(op->get_shape(), (Shape{2, 96, 17, 31}));
 }
+
+TEST(type_prop_layers, psroi_pooling)
+{
+    auto inputs = make_shared<op::Parameter>(element::f32, Shape{1, 3, 4, 5});
+    auto coords = make_shared<op::Parameter>(element::f32, Shape{150, 5});
+    auto op = make_shared<op::PSROIPooling>(inputs, coords, 2, 6, 0.0625, Shape{}, "Avg");
+    ASSERT_EQ(op->get_shape(), (Shape{150, 2, 6, 6}));
+}
+
+TEST(type_prop_layers, roi_pooling)
+{
+    auto inputs = make_shared<op::Parameter>(element::f32, Shape{2, 3, 4, 5});
+    auto coords = make_shared<op::Parameter>(element::f32, Shape{150, 5});
+    auto op = make_shared<op::ROIPooling>(inputs, coords, Shape{6, 6}, 0.0625, "Max");
+    ASSERT_EQ(op->get_shape(), (Shape{150, 3, 6, 6}));
+}