MaxPool backprop requires knowledge of y. :(

Author: staticfloat

src/impl/pooling_direct.jl

@@ -123,8 +123,8 @@ for name in (:max, :mean)
     # it's unfortunately different enough that I think we need a separate function.  :(
     @eval @timeit_debug to function $((Symbol("∇$(name)pool_direct!")))(
                     dx::AbstractArray{T,5}, dy::AbstractArray{T,5},
-                    x::AbstractArray{T,5}, pdims::PoolDims;
-                    alpha::T = T(1), beta::T = T(0)) where {T}
+                    y::AbstractArray{T,5}, x::AbstractArray{T,5},
+                    pdims::PoolDims; alpha::T = T(1), beta::T = T(0)) where {T}
         check_dims(size(x), size(dy), pdims)
 
         width, height, depth = input_size(pdims)
@@ -156,7 +156,8 @@ for name in (:max, :mean)
             h in h_region,
             w in w_region
 
-            # Grab the incoming gradient at this index for future use
+            # Grab the output at this index for future use
+            y_idx = y[w, h, d, c, batch_idx]
             dy_idx = dy[w, h, d, c, batch_idx]
             maxpool_already_chose = false
 
@@ -174,7 +175,7 @@ for name in (:max, :mean)
                 if $(name == :max)
                     # If it's equal; this is the one we chose. We only choose one per
                     # kernel window, all other elements of dx must be zero.
-                    if dy_idx == x[x_idxs...] && !maxpool_already_chose
+                    if y_idx == x[x_idxs...] && !maxpool_already_chose
                         dx[x_idxs...] = dy_idx*alpha + beta*dx[x_idxs...]
                         maxpool_already_chose = true
                     # Maxpooling does not support `beta` right now.  :(
@@ -199,6 +200,7 @@ for name in (:max, :mean)
                 w in w_region
 
                 # Grab the incoming gradient at this index for future use
+                y_idx = dy[w, h, d, c, batch_idx]
                 dy_idx = dy[w, h, d, c, batch_idx]
                 maxpool_already_chose = false
 
@@ -225,7 +227,7 @@ for name in (:max, :mean)
                             # Same as above
                             x_idxs = (input_kw, input_kh, input_kd, c, batch_idx)
                             if $(name == :max)
-                                if dy_idx == x[x_idxs...] && !maxpool_already_chose
+                                if y_idx == x[x_idxs...] && !maxpool_already_chose
                                     dx[x_idxs...] = dy_idx*alpha + beta*dx[x_idxs...]
                                     maxpool_already_chose = true
                                 #else

src/pooling.jl

@@ -47,9 +47,10 @@ for (front_name, backend) in (
     )
     @eval begin
         function $(Symbol("$(front_name)!"))(
-                dx::AbstractArray{T,5}, dy::AbstractArray{T,5},
-                x::AbstractArray{T,5}, pdims::PoolDims; kwargs...) where {T}
-            $(Symbol("$(front_name)_$(backend)!"))(dx, dy, x, pdims; kwargs...)
+                        dx::AbstractArray{T,5}, dy::AbstractArray{T,5},
+                        y::AbstractArray{T,5}, x::AbstractArray{T,5},
+                        pdims::PoolDims; kwargs...) where {T}
+            $(Symbol("$(front_name)_$(backend)!"))(dx, dy, y, x, pdims; kwargs...)
         end
     end
 end
@@ -79,12 +80,13 @@ for front_name in (:maxpool, :meanpool)
 
                 # backprops too
                 function $(Symbol("∇$(front_name)$(backend)!"))(
-                        dx::AbstractArray{T,$N}, dy::AbstractArray{T,$N},
-                        x::AbstractArray{T,$N}, pdims::PoolDims;
-                        kwargs...) where {T}
+                                dx::AbstractArray{T,$N}, dy::AbstractArray{T,$N},
+                                y::AbstractArray{T,$N}, x::AbstractArray{T,$N},
+                                pdims::PoolDims; kwargs...) where {T}
                     $(Symbol("∇$(front_name)$(backend)!"))(
                         insert_singleton_spatial_dimension(dx, $(5 - N)),
                         insert_singleton_spatial_dimension(dy, $(5 - N)),
+                        insert_singleton_spatial_dimension(y, $(5 - N)),
                         insert_singleton_spatial_dimension(x, $(5 - N)),
                         insert_singleton_spatial_dimension(pdims, $(5 - N));
                         kwargs...
@@ -114,10 +116,11 @@ for backend in (Symbol(), :_direct, :_im2col)
 
             # Backprops too
             @timeit_debug to function $(Symbol("∇$(name)$(backend)"))(
-                            dy::AbstractArray{T,N}, x::AbstractArray{T},
-                            pdims::PoolDims; kwargs...) where {T, N}
+                            dy::AbstractArray{T,N}, y::AbstractArray{T,N},
+                            x::AbstractArray{T,N}, pdims::PoolDims;
+                            kwargs...) where {T, N}
                 dx = zeros(T, input_size(pdims)..., channels_in(pdims), size(dy, N))
-                return $(Symbol("∇$(name)$(backend)!"))(dx, dy, x, pdims; kwargs...)
+                return $(Symbol("∇$(name)$(backend)!"))(dx, dy, y, x, pdims; kwargs...)
             end
         end
     end

test/pooling.jl

@@ -280,19 +280,19 @@ for rank in (1, 2, 3)
                 pdims = PoolDims(x, 2)
                 y_hat = pool(x, pdims)
                 @test ddims(y_hat) == y
-                @test ddims(∇pool(y_hat, x, pdims)) == dx
+                @test ddims(∇pool(y_hat, y_hat, x, pdims)) == dx
 
                 # Strided pooling
                 pdims = PoolDims(x, 2; stride=1)
                 y_hat = pool(x, pdims)
                 @test ddims(y_hat) == y_nostride
-                @test ddims(∇pool(y_hat, x, pdims)) == dx_nostride
+                @test ddims(∇pool(y_hat, y_hat, x, pdims)) == dx_nostride
 
                 # Padded pooling
                 pdims = PoolDims(x, 2; padding=1)
                 y_hat = pool(x, pdims)
                 @test ddims(y_hat) == y_pad
-                @test ddims(∇pool(y_hat, x, pdims)) == dx_pad
+                @test ddims(∇pool(y_hat, y_hat, x, pdims)) == dx_pad
             end
         end
     end