Commit 277b5dcb authored by Mateusz Bencer's avatar Mateusz Bencer Committed by Scott Cyphers

PSROI Pooling was adjusted (#3347)

parent d116e47e
......@@ -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);
}
......@@ -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;
};
}
}
......@@ -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}));
}
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment