Add padding

Author: americast

src/pool.jl

@@ -1,15 +1,15 @@
 import CUDAnative
 
-function maxpool2d_kernel(state, A::AbstractArray{T}, out, Asize, pool, stride_, outSize) where T
+function maxpool2d_kernel(state, A::AbstractArray{T}, out, Asize, pool, stride, outSize) where T
     ilin = linear_index(state)
     idx = GPUArrays.gpu_ind2sub(Asize, ilin)
     if (idx[1] > outSize[1] || idx[2] > outSize[2] || idx[3] > outSize[3] || idx[4] > outSize[4])
         return
     end
 
-    temp_max = A[((idx[1] - 1) * stride_) + Asize[1] * (idx[2] - 1) * stride_ + (Asize[1] * Asize[2]) * (idx[3] - 1) + (Asize[1] * Asize[2] * Asize[3]) * (idx[4] - 1) + 1]
-    max_pos = ((idx[1] - 1) * stride_) + Asize[1] * (idx[2] - 1) * stride_ + (Asize[1] * Asize[2]) * (idx[3] - 1) + (Asize[1] * Asize[2] * Asize[3]) * (idx[4] - 1) + 1
-    curr_pos = ((idx[1] - 1) * stride_) + Asize[1] * (idx[2] - 1) * stride_ + (Asize[1] * Asize[2]) * (idx[3] - 1) + (Asize[1] * Asize[2] * Asize[3]) * (idx[4] - 1) + 1
+    temp_max = A[((idx[1] - 1) * stride) + Asize[1] * (idx[2] - 1) * stride + (Asize[1] * Asize[2]) * (idx[3] - 1) + (Asize[1] * Asize[2] * Asize[3]) * (idx[4] - 1) + 1]
+    max_pos = ((idx[1] - 1) * stride) + Asize[1] * (idx[2] - 1) * stride + (Asize[1] * Asize[2]) * (idx[3] - 1) + (Asize[1] * Asize[2] * Asize[3]) * (idx[4] - 1) + 1
+    curr_pos = ((idx[1] - 1) * stride) + Asize[1] * (idx[2] - 1) * stride + (Asize[1] * Asize[2]) * (idx[3] - 1) + (Asize[1] * Asize[2] * Asize[3]) * (idx[4] - 1) + 1
 
     for p in 1:pool
         for p in 1:pool
@@ -27,14 +27,16 @@ function maxpool2d_kernel(state, A::AbstractArray{T}, out, Asize, pool, stride_,
 end
 
 
-function maxpool2d(a, pool; stride_ = 1)
-    Asize = UInt32.(size(a))
+function maxpool2d(a, pool; stride = 1, pad = 0)
+    b = zeros(typeof(a), size(a,1) + pad * 2, size(a,2) + pad * 2, size(a,3), size(a,4))
+    b[pad + 1 : pad + size(a,1), pad + 1 : pad + size(a,2), :, :] = a
+    Asize = UInt32.(size(b))
     pool = UInt32(pool)
-    stride_ = UInt32(stride_)
-    out = similar(a)
-    out = out[1:(div(Asize[1] - pool, stride_) + 1), 1:(div(Asize[2] - pool, stride_) + 1), :, :]
+    stride = UInt32(stride)
+    out = similar(b)
+    out = out[1:(div(Asize[1] - pool, stride) + 1), 1:(div(Asize[2] - pool, stride) + 1), :, :]
     outSize = UInt32.(size(out))
-    gpu_call(maxpool2d_kernel, a, (a, out, Asize, pool, stride_, outSize))
+    gpu_call(maxpool2d_kernel, b, (b, out, Asize, pool, stride, outSize))
     GPUArrays.synchronize(out)
     out
 end