channel-pooling along any dimension for ND-tensor (CPU)

Author: pluskid

src/cuda/layers/channel-pooling.jl

@@ -1,16 +1,14 @@
-function setup_etc(backend::GPUBackend, layer::ChannelPoolingLayer, inputs, pooled_chann)
+function setup_etc(backend::GPUBackend, layer::ChannelPoolingLayer, inputs, blobs)
   if isa(layer.pooling, Pooling.Max)
     masks = Array(CuPtr, length(inputs))
     for i = 1:length(inputs)
-      masks[i] = CUDA.cualloc(Csize_t, get_width(inputs[i]) * get_height(inputs[i]) *
-          pooled_chann[i] * get_num(inputs[i]))
+      masks[i] = CUDA.cualloc(Csize_t, length(blobs[i]))
     end
     etc = masks
   elseif isa(layer.pooling, Pooling.Mean)
     integrals = Array(CuPtr, length(inputs))
     for i = 1:length(inputs)
-      integrals[i] = CUDA.cualloc(eltype(inputs[i]), get_width(inputs[i]) * get_height(inputs[i]) *
-          get_chann(inputs[i]))
+      integrals[i] = CUDA.cualloc(eltype(inputs[i]), prod(size(inputs[i])[1:end-1]))
     end
     etc = integrals
   else

src/layers/channel-pooling.jl

@@ -21,19 +21,17 @@ type ChannelPoolingLayerState <: LayerState
   etc        :: Any
 end
 
-function setup_etc(backend::CPUBackend, layer::ChannelPoolingLayer, inputs, pooled_chann)
+function setup_etc(backend::CPUBackend, layer::ChannelPoolingLayer, inputs, blobs)
   if isa(layer.pooling, Pooling.Max)
     masks = Array(Array, length(inputs))
     for i = 1:length(inputs)
-      masks[i] = Array(Csize_t, get_width(inputs[i]), get_height(inputs[i]),
-          pooled_chann[i], get_num(inputs[i]))
+      masks[i] = Array(Csize_t, size(blobs[i]))
     end
     etc = masks
   elseif isa(layer.pooling, Pooling.Mean)
     integrals = Array(Array, length(inputs))
     for i = 1:length(inputs)
-      integrals[i] = Array(eltype(inputs[i]), get_width(inputs[i]), get_height(inputs[i]),
-          get_chann(inputs[i]))
+      integrals[i] = Array(eltype(inputs[i]), size(inputs[i])[1:end-1])
     end
     etc = integrals
   else
@@ -43,40 +41,43 @@ function setup_etc(backend::CPUBackend, layer::ChannelPoolingLayer, inputs, pool
 end
 
 function setup(backend::Backend, layer::ChannelPoolingLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
-  for i = 1:length(inputs)
-    # currently we only handle 4D-tensor
-    @assert ndims(inputs[i]) == 4
-  end
-
   pooled_chann_all = Array(Int, length(inputs))
   blobs = Array(Blob, length(inputs))
   blobs_diff = Array(Blob, length(inputs))
   op_dims = Array(Int, length(inputs))
 
   for i = 1:length(inputs)
     dim_total = ndims(inputs[i])
-    dim = layer.dim < 0 ? layer.dim + dim_total+1 : layer.dim
-    op_dims[i] = dim
+    op_dim = layer.dim < 0 ? layer.dim + dim_total+1 : layer.dim
+    @assert 1 <= op_dim <= dim_total
+    @assert op_dim != dim_total
 
-    width, height, channels, num = size(inputs[i])
-    pooled_chann = int(ceil(float(channels + layer.pad[1]+layer.pad[2] - layer.kernel) / layer.stride)) + 1
+    op_dims[i] = op_dim
+
+    dims = [size(inputs[i])...]
+    pool_dim = dims[op_dim]
+    pooled_dim = int(ceil(float(pool_dim + layer.pad[1]+layer.pad[2] - layer.kernel) / layer.stride)) + 1
 
     # make sure the last pooling is not purely pooling padded area
-    if ((pooled_chann-1)*layer.stride >= channels + layer.pad[1])
-      pooled_chann -= 1
+    if ((pooled_dim-1)*layer.stride >= pool_dim + layer.pad[1])
+      pooled_dim -= 1
     end
-    pooled_chann_all[i] = pooled_chann
+    pooled_chann_all[i] = pooled_dim
+
+    output_dims = copy(dims)
+    output_dims[op_dim] = pooled_dim
+    output_dims = tuple(output_dims...)
 
     data_type = eltype(inputs[i])
-    blobs[i] = make_blob(backend, data_type, (width,height,pooled_chann_all[i],num))
+    blobs[i] = make_blob(backend, data_type, output_dims)
     if isa(diffs[i], NullBlob)
       blobs_diff[i] = NullBlob()
     else
-      blobs_diff[i] = make_blob(backend, data_type, (width,height,pooled_chann_all[i],num))
+      blobs_diff[i] = make_blob(backend, data_type, output_dims)
     end
   end
 
-  etc = setup_etc(backend, layer, inputs, pooled_chann_all)
+  etc = setup_etc(backend, layer, inputs, blobs)
   state = ChannelPoolingLayerState(layer, blobs, blobs_diff, op_dims, etc)
 end
 function shutdown_etc(backend::CPUBackend, state::ChannelPoolingLayerState)

test/layers/channel-pooling.jl

@@ -1,15 +1,16 @@
 function test_channel_pooling_layer(backend::Backend, pooling::PoolingFunction, tensor_dim::Int, n_input, T, eps)
   println("-- Testing ChannelPooling($(typeof(pooling))) on $(typeof(backend)){$T}...")
-  println("    > Setup")
 
   dims = [abs(rand(Int, tensor_dim)) % 7 + 7 for i = 1:n_input]
-  op_dim = 3
+  op_dim = max(abs(rand(Int)) % tensor_dim, 1)
   pad = (2,2)
   kernel = 3
   stride = 2
 
+  println("    > Setup (pool along dimension $op_dim for $tensor_dim-D tensors)")
+
   layer = ChannelPoolingLayer(kernel=kernel, stride=stride, pad=pad, pooling=pooling,
-      tops=Array(Symbol,n_input), bottoms=Array(Symbol,n_input))
+      tops=Array(Symbol,n_input), bottoms=Array(Symbol,n_input), dim=op_dim)
 
   input = [rand(T, dim...) for dim in dims]
   inputs = Blob[make_blob(backend, x) for x in input]
@@ -118,7 +119,7 @@ function channel_pooling_backward(state, i, input::Array, diff::Array, payload::
 end
 
 function test_channel_pooling_layer(backend::Backend, pooling::PoolingFunction, n_input, T, eps)
-  for i = 4:4
+  for i = 2:6
     test_channel_pooling_layer(backend, pooling, i, n_input, T, eps)
   end
 end