From 760ffaa51475e95c9e64ec858a7d0c94fd7937ae Mon Sep 17 00:00:00 2001
From: Sergio <sguada@gmail.com>
Date: Fri, 3 Oct 2014 08:19:40 -0700
Subject: [PATCH] Added global_pooling to set the kernel size equal to the
 bottom size

Added check for padding and stride with global_pooling
---
 include/caffe/vision_layers.hpp       |  1 +
 src/caffe/layers/pooling_layer.cpp    | 32 ++++++++++++++++++++++-----
 src/caffe/proto/caffe.proto           |  3 +++
 src/caffe/test/test_pooling_layer.cpp | 14 ++++++++++++
 4 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/include/caffe/vision_layers.hpp b/include/caffe/vision_layers.hpp
index 58ac5344da4..c803cd72449 100644
--- a/include/caffe/vision_layers.hpp
+++ b/include/caffe/vision_layers.hpp
@@ -319,6 +319,7 @@ class PoolingLayer : public Layer<Dtype> {
   int channels_;
   int height_, width_;
   int pooled_height_, pooled_width_;
+  bool global_pooling_;
   Blob<Dtype> rand_idx_;
   Blob<int> max_idx_;
 };
diff --git a/src/caffe/layers/pooling_layer.cpp b/src/caffe/layers/pooling_layer.cpp
index 26c92c13b63..2bfbb01fcec 100644
--- a/src/caffe/layers/pooling_layer.cpp
+++ b/src/caffe/layers/pooling_layer.cpp
@@ -17,12 +17,18 @@ template <typename Dtype>
 void PoolingLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
       const vector<Blob<Dtype>*>& top) {
   PoolingParameter pool_param = this->layer_param_.pooling_param();
-  CHECK(!pool_param.has_kernel_size() !=
+  if (pool_param.global_pooling()) {
+    CHECK(!(pool_param.has_kernel_size() ||
+      pool_param.has_kernel_h() || pool_param.has_kernel_w()))
+      << "With Global_pooling: true Filter size cannot specified";
+  } else {
+    CHECK(!pool_param.has_kernel_size() !=
       !(pool_param.has_kernel_h() && pool_param.has_kernel_w()))
       << "Filter size is kernel_size OR kernel_h and kernel_w; not both";
-  CHECK(pool_param.has_kernel_size() ||
+    CHECK(pool_param.has_kernel_size() ||
       (pool_param.has_kernel_h() && pool_param.has_kernel_w()))
       << "For non-square filters both kernel_h and kernel_w are required.";
+  }
   CHECK((!pool_param.has_pad() && pool_param.has_pad_h()
       && pool_param.has_pad_w())
       || (!pool_param.has_pad_h() && !pool_param.has_pad_w()))
@@ -31,11 +37,17 @@ void PoolingLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
       && pool_param.has_stride_w())
       || (!pool_param.has_stride_h() && !pool_param.has_stride_w()))
       << "Stride is stride OR stride_h and stride_w are required.";
-  if (pool_param.has_kernel_size()) {
-    kernel_h_ = kernel_w_ = pool_param.kernel_size();
+  global_pooling_ = pool_param.global_pooling();
+  if (global_pooling_) {
+    kernel_h_ = bottom[0]->height();
+    kernel_w_ = bottom[0]->width();
   } else {
-    kernel_h_ = pool_param.kernel_h();
-    kernel_w_ = pool_param.kernel_w();
+    if (pool_param.has_kernel_size()) {
+      kernel_h_ = kernel_w_ = pool_param.kernel_size();
+    } else {
+      kernel_h_ = pool_param.kernel_h();
+      kernel_w_ = pool_param.kernel_w();
+    }
   }
   CHECK_GT(kernel_h_, 0) << "Filter dimensions cannot be zero.";
   CHECK_GT(kernel_w_, 0) << "Filter dimensions cannot be zero.";
@@ -51,6 +63,10 @@ void PoolingLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
     stride_h_ = pool_param.stride_h();
     stride_w_ = pool_param.stride_w();
   }
+  if (global_pooling_) {
+    CHECK(pad_h_ == 0 && pad_w_ == 0 && stride_h_ == 1 && stride_w_ == 1)
+      << "With Global_pooling: true; only pad = 0 and stride = 1";
+  }
   if (pad_h_ != 0 || pad_w_ != 0) {
     CHECK(this->layer_param_.pooling_param().pool()
         == PoolingParameter_PoolMethod_AVE
@@ -68,6 +84,10 @@ void PoolingLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
   channels_ = bottom[0]->channels();
   height_ = bottom[0]->height();
   width_ = bottom[0]->width();
+  if (global_pooling_) {
+    kernel_h_ = bottom[0]->height();
+    kernel_w_ = bottom[0]->width();
+  }
   pooled_height_ = static_cast<int>(ceil(static_cast<float>(
       height_ + 2 * pad_h_ - kernel_h_) / stride_h_)) + 1;
   pooled_width_ = static_cast<int>(ceil(static_cast<float>(
diff --git a/src/caffe/proto/caffe.proto b/src/caffe/proto/caffe.proto
index 6944ae847ff..9374abe5b70 100644
--- a/src/caffe/proto/caffe.proto
+++ b/src/caffe/proto/caffe.proto
@@ -602,6 +602,9 @@ message PoolingParameter {
     CUDNN = 2;
   }
   optional Engine engine = 11 [default = DEFAULT];
+  // If global_pooling then it will pool over the size of the bottom by doing
+  // kernel_h = bottom->height and kernel_w = bottom->width
+  optional bool global_pooling = 12 [default = false];
 }
 
 // Message that stores parameters used by PowerLayer
diff --git a/src/caffe/test/test_pooling_layer.cpp b/src/caffe/test/test_pooling_layer.cpp
index f52073c4e6f..435caa8381e 100644
--- a/src/caffe/test/test_pooling_layer.cpp
+++ b/src/caffe/test/test_pooling_layer.cpp
@@ -400,6 +400,20 @@ TYPED_TEST(PoolingLayerTest, TestSetupPadded) {
   EXPECT_EQ(this->blob_top_->width(), 3);
 }
 
+TYPED_TEST(PoolingLayerTest, TestSetupGlobalPooling) {
+  typedef typename TypeParam::Dtype Dtype;
+  LayerParameter layer_param;
+  PoolingParameter* pooling_param = layer_param.mutable_pooling_param();
+  pooling_param->set_global_pooling(true);
+  pooling_param->set_pool(PoolingParameter_PoolMethod_AVE);
+  PoolingLayer<Dtype> layer(layer_param);
+  layer.SetUp(this->blob_bottom_vec_, this->blob_top_vec_);
+  EXPECT_EQ(this->blob_top_->num(), this->blob_bottom_->num());
+  EXPECT_EQ(this->blob_top_->channels(), this->blob_bottom_->channels());
+  EXPECT_EQ(this->blob_top_->height(), 1);
+  EXPECT_EQ(this->blob_top_->width(), 1);
+}
+
 /*
 TYPED_TEST(PoolingLayerTest, PrintBackward) {
   LayerParameter layer_param;