PSROI Pooling was adjusted (#3347)

src/ngraph/op/experimental/layers/psroi_pooling.cpp

@@ -22,18 +22,20 @@ using namespace ngraph;
 const string op::PSROIPooling::type_name{"PSROIPooling"};
 
 op::PSROIPooling::PSROIPooling(const shared_ptr<Node>& input,
-                               const std::shared_ptr<Node>& coords,
+                               const 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)
+                               int spatial_bins_x,
+                               int spatial_bins_y,
+                               const string& mode)
     : Op(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)
+    , m_spatial_bins_x(spatial_bins_x)
+    , m_spatial_bins_y(spatial_bins_y)
+    , m_mode(mode)
 {
     constructor_validate_and_infer_types();
 }
@@ -74,6 +76,7 @@ shared_ptr<Node> op::PSROIPooling::copy_with_new_args(const NodeVector& new_args
                                      m_output_dim,
                                      m_group_size,
                                      m_spatial_scale,
-                                     m_num_bins,
-                                     m_kind);
+                                     m_spatial_bins_x,
+                                     m_spatial_bins_y,
+                                     m_mode);
 }

src/ngraph/op/experimental/layers/psroi_pooling.hpp

@@ -35,15 +35,17 @@ namespace ngraph
             /// \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
+            /// \param spatial_bins_x Numbers of bins to divide the input feature maps over width
+            /// \param spatial_bins_y Numbers of bins to divide the input feature maps over height
+            /// \param mode           Mode 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);
+                         int spatial_bins_x,
+                         int spatial_bins_y,
+                         const std::string& mode);
 
             void validate_and_infer_types() override;
 
@@ -53,14 +55,16 @@ namespace ngraph
             size_t get_output_dim() const { return m_output_dim; }
             size_t get_group_size() const { return m_group_size; }
             float get_spatial_scale() const { return m_spatial_scale; }
-            const Shape& get_num_bins() const { return m_num_bins; }
-            const std::string& get_kind() const { return m_kind; }
+            int get_spatial_bins_x() const { return m_spatial_bins_x; }
+            int get_spatial_bins_y() const { return m_spatial_bins_y; }
+            const std::string& get_mode() const { return m_mode; }
         private:
             size_t m_output_dim;
             size_t m_group_size;
             float m_spatial_scale;
-            Shape m_num_bins;
-            std::string m_kind;
+            int m_spatial_bins_x;
+            int m_spatial_bins_y;
+            std::string m_mode;
         };
     }
 }

test/type_prop_layers.cpp

@@ -171,7 +171,7 @@ 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");
+    auto op = make_shared<op::PSROIPooling>(inputs, coords, 2, 6, 0.0625, 0, 0, "Avg");
     ASSERT_EQ(op->get_shape(), (Shape{150, 2, 6, 6}));
 }