From 822dabc37ae497e6f6eb48498e86489e33189ba7 Mon Sep 17 00:00:00 2001
From: Chiyuan Zhang <pluskid@gmail.com>
Date: Fri, 19 Dec 2014 23:28:11 +0800
Subject: [PATCH] channel pooling along any dimension (GPU)

---
 src/cuda/layers/channel-pooling.jl | 36 +++++++++++++-----------------
 src/layers/channel-pooling.jl      |  4 ----
 2 files changed, 16 insertions(+), 24 deletions(-)

diff --git a/src/cuda/layers/channel-pooling.jl b/src/cuda/layers/channel-pooling.jl
index f6ae60a..073fbaf 100644
--- a/src/cuda/layers/channel-pooling.jl
+++ b/src/cuda/layers/channel-pooling.jl
@@ -37,9 +37,9 @@ function forward(backend::GPUBackend, pool::StdPoolingFunction,
     output = state.blobs[i]
 
     if isa(pool, Pooling.Max)
-      cuda_max_channel_pooling_forward(backend, input, output, state.etc[i], state.layer)
+      cuda_max_channel_pooling_forward(backend, input, output, state.etc[i], state.layer, state.op_dims[i])
     elseif isa(pool, Pooling.Mean)
-      cuda_mean_channel_pooling_forward(backend, input, output, state.etc[i], state.layer)
+      cuda_mean_channel_pooling_forward(backend, input, output, state.etc[i], state.layer, state.op_dims[i])
     else
       error("Pooling for $pool not implemented yet")
     end
@@ -57,9 +57,9 @@ function backward(backend::GPUBackend, pool::StdPoolingFunction, state::ChannelP
     diff = diffs[i]
     if !isa(diff, NullBlob)
       if isa(pool, Pooling.Max)
-        cuda_max_channel_pooling_backward(backend, diff, state.blobs_diff[i], state.etc[i], state.layer)
+        cuda_max_channel_pooling_backward(backend, diff, state.blobs_diff[i], state.etc[i], state.layer, state.op_dims[i])
       elseif isa(pool, Pooling.Mean)
-        cuda_mean_channel_pooling_backward(backend, diff, state.blobs_diff[i], state.layer)
+        cuda_mean_channel_pooling_backward(backend, diff, state.blobs_diff[i], state.layer, state.op_dims[i])
       else
         error("Pooling for $pool not implemented yet")
       end
@@ -70,15 +70,14 @@ function backward(backend::GPUBackend, pool::StdPoolingFunction, state::ChannelP
 end
 
 function cuda_mean_channel_pooling_forward{T}(backend::GPUBackend, input::CuTensorBlob{T},
-    output::CuTensorBlob{T}, integral::CuPtr, layer)
+    output::CuTensorBlob{T}, integral::CuPtr, layer, op_dim)
 
-  width, height, channels, num = size(input)
-  pooled_chann = size(output, 3)
+  spatial_dim_T, channels, num = split_dims(input, op_dim)
+  pooled_chann = size(output, op_dim)
   one = convert(T, 1)
   neg_one = convert(T, -1)
   scale = convert(T, 1.0/layer.kernel)
 
-  spatial_dim_T = width*height
   spatial_dim = spatial_dim_T * sizeof(T)
   fea_dim = spatial_dim * channels
   output_fea_dim = spatial_dim * pooled_chann
@@ -116,15 +115,14 @@ function cuda_mean_channel_pooling_forward{T}(backend::GPUBackend, input::CuTens
 end
 
 function cuda_mean_channel_pooling_backward{T}(backend::GPUBackend, input::CuTensorBlob{T},
-    output::CuTensorBlob{T}, layer)
+    output::CuTensorBlob{T}, layer, op_dim)
 
-  width, height, channels, num = size(input)
-  pooled_chann = size(output, 3)
+  spatial_dim_T, channels, num = split_dims(input, op_dim)
+  pooled_chann = size(output, op_dim)
   scale = 1/convert(T, layer.kernel)
 
   fill!(input, 0)
 
-  spatial_dim_T = width*height
   spatial_dim = spatial_dim_T * sizeof(T)
   fea_dim = spatial_dim * channels
   output_fea_dim = spatial_dim * pooled_chann
@@ -157,11 +155,10 @@ function cuda_geometry_max_chann_pool(sp_dim::Int, num::Int)
 
 end
 function cuda_max_channel_pooling_forward{T}(backend::GPUBackend, input::CuTensorBlob{T},
-    output::CuTensorBlob{T}, mask::CuPtr, layer)
+    output::CuTensorBlob{T}, mask::CuPtr, layer, op_dim)
 
-  width, height, channels, num = size(input)
-  sp_dim = width*height
-  pooled_chann = get_chann(output)
+  sp_dim, channels, num = split_dims(input, op_dim)
+  pooled_chann = size(output, op_dim)
 
   cuda_dim = cuda_geometry_max_chann_pool(sp_dim, num);
   if T == Float32
@@ -177,11 +174,10 @@ function cuda_max_channel_pooling_forward{T}(backend::GPUBackend, input::CuTenso
 end
 
 function cuda_max_channel_pooling_backward{T}(backend::GPUBackend, input::CuTensorBlob{T},
-    output::CuTensorBlob{T}, mask::CuPtr, layer)
+    output::CuTensorBlob{T}, mask::CuPtr, layer, op_dim)
 
-  width, height, channels, num = size(input)
-  sp_dim = width*height
-  pooled_chann = get_chann(output)
+  sp_dim, channels, num = split_dims(input, op_dim)
+  pooled_chann = size(output, op_dim)
 
   cuda_dim = cuda_geometry_max_chann_pool(sp_dim, num);
   if T == Float32
diff --git a/src/layers/channel-pooling.jl b/src/layers/channel-pooling.jl
index cc85f5a..d05751b 100644
--- a/src/layers/channel-pooling.jl
+++ b/src/layers/channel-pooling.jl
@@ -105,8 +105,6 @@ function forward(backend::CPUBackend, pool::StdPoolingFunction,
       max_channel_pooling_forward(reshape(input,dims_in), reshape(output,dims_out), reshape(state.etc[i],dims_out), state.layer)
     elseif isa(pool, Pooling.Mean)
       mean_channel_pooling_forward(reshape(input,dims_in), reshape(output,dims_out), state.etc[i], state.layer)
-    else
-      error("Pooling for $pool not implemented yet")
     end
   end
 end
@@ -130,8 +128,6 @@ function backward(backend::CPUBackend, pool::StdPoolingFunction, state::ChannelP
             reshape(state.etc[i],dims_out), state.layer)
       elseif isa(pool, Pooling.Mean)
         mean_channel_pooling_backward(reshape(diff.data,dims_in), reshape(state.blobs_diff[i].data,dims_out), state.layer)
-      else
-        error("Pooling for $pool not implemented yet")
       end
     end
   end