Add lppool implementation (#447)

* add: init my lppool impl

* add: impl gradient calculate

* test: for lppool(1d & 2d)

* doc: juliadoc & formula comment

* doc: update for meanpool/maxpool, add params details and warning for lppool-maxpool

* update: 1 => T(1); remove unnecessary if-condition; clear error msg and doc

* update: move dim checker from exported pool function to PoolDims

* rename: lppool => normpool

* add: normpool p value checker, p must be in (0, Inf)

* rename: normpool => lpnormpool

* doc: add paper reference to lp pooling

* doc: usage and parameter description of pool functions

Author: skyleaworlder

docs/src/reference.md

@@ -44,12 +44,13 @@ logsoftmax
 
 ## Pooling
 
-`Flux`'s `AdaptiveMaxPool`, `AdaptiveMeanPool`, `GlobalMaxPool`, `GlobalMeanPool`, `MaxPool`, and `MeanPool` use `NNlib.PoolDims`, `NNlib.maxpool`, and `NNlib.meanpool` as their backend.
+`Flux`'s `AdaptiveMaxPool`, `AdaptiveMeanPool`, `GlobalMaxPool`, `GlobalMeanPool`, `MaxPool`, `MeanPool` and `lpnormpool` use `NNlib.PoolDims`, `NNlib.maxpool`, `NNlib.meanpool` and `NNlib.lpnormpool` as their backend.
 
 ```@docs
 PoolDims
 maxpool
 meanpool
+lpnormpool
 ```
 
 ## Padding

src/NNlib.jl

@@ -71,8 +71,8 @@ include("ctc.jl")
 export ctc_loss
 
 include("pooling.jl")
-export maxpool, maxpool!, meanpool, meanpool!, 
-    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!
+export maxpool, maxpool!, meanpool, meanpool!, lpnormpool, lpnormpool!,
+    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!, ∇lpnormpool, ∇lpnormpool!
 
 include("padding.jl")
 export pad_constant, pad_repeat, pad_reflect, pad_zeros

src/dim_helpers/PoolDims.jl

@@ -25,6 +25,12 @@ function PoolDims(
     _check_kernel(k::NTuple, ::Int) = k
 
     kernel = _check_kernel(k, M - 2)
+    length(x_size) == length(kernel) + 2 || error(
+        "PoolDims expects ndim(x) == length(k)+2 or length(size(x)) == length(kernel)+2,
+        dimension of x_size is $(length(x_size)),
+        length of k need $(length(x_size) - 2),
+        but now it's $(length(kernel))"
+    )
     spdf_kernel = NTuple{M, Int}([kernel..., 1, 1])
 
     sstride, ppadding, ddilation = check_spdf(

src/impl/pooling_direct.jl

@@ -1,14 +1,14 @@
 # Pooling is so similar, we abstract over meanpooling and maxpooling, simply replacing
 # the inner loop operation and a few initialization parameters.
-for name in (:max, :mean)
+for name in (:max, :mean, :lpnorm)
     @eval function $((Symbol("$(name)pool_direct!")))(
                     y::AbstractArray{T, 5}, x::AbstractArray{T, 5},
-                    pdims::PoolDims; alpha::T=T(1), beta::T=T(0)) where T
+                    pdims::PoolDims; alpha::T=T(1), beta::T=T(0), kwargs...) where T
         $((Symbol("$(name)pool_direct!")))(
             y, x, pdims,
             Val(kernel_size(pdims)), Val(channels_out(pdims)),
             Val(padding(pdims)), Val(dilation(pdims)), Val(stride(pdims));
-            alpha, beta)
+            alpha, beta, kwargs...)
         return y
     end
 
@@ -17,7 +17,7 @@ for name in (:max, :mean)
         pdims::PoolDims,
         # kernel size, channels out, padding, dilation, stride
         ::Val{K}, ::Val{C}, ::Val{P}, ::Val{D}, ::Val{S};
-        alpha::T=T(1), beta::T=T(0),
+        alpha::T=T(1), beta::T=T(0), kwargs...
     ) where {T, K, C, P, D, S}
         @assert beta == T(0) "beta not supported yet"
         check_dims(size(x), size(y), pdims)
@@ -41,10 +41,15 @@ for name in (:max, :mean)
             alpha = alpha / prod(K)
         end
 
+        p = if $(name != :lpnorm) 0 else
+            !haskey(kwargs, :p) && error("lpnormpool needs keyword argument `p`")
+            kwargs[:p]
+        end
+
         # Each loop, we initialize `m` to something, set that here.
         m_init = if $(name == :max)
             T <: AbstractFloat ? nextfloat(typemin(T)) : typemin(T)
-        elseif $(name == :mean)
+        elseif $(name == :mean) || $(name == :lpnorm)
             T(0)
         else
             error("Unimplemented codegen path")
@@ -78,11 +83,17 @@ for name in (:max, :mean)
                     end
                 elseif $(name == :mean)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]
+                elseif $(name == :lpnorm)
+                    # y = (∑ᵢ xᵢ^p)^(1 / p), here to calculate ∑ᵢ xᵢ^p
+                    m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                 else
                     error("Unimplemented codegen path")
                 end
             end
 
+            # for lpnormpool, y = (∑ᵢ xᵢ^p)^(1 / p)
+            m = $(name == :lpnorm) ? m^(T(1) / p) : m
+
             y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
             end
             end
@@ -128,12 +139,15 @@ for name in (:max, :mean)
                                 end
                             elseif $(name == :mean)
                                 m += x[input_kw, input_kh, input_kd, c, batch_idx]
+                            elseif $(name == :lpnorm)
+                                m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                             else
                                 error("Unimplemented codegen path")
                             end
                         end
                     end
                 end
+                $(name == :lpnorm) && (m = m^(T(1) / p))
                 y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
                 end
                 end
@@ -159,7 +173,7 @@ for name in (:max, :mean)
                     dx::AbstractArray{T,5}, dy::AbstractArray{T,5},
                     y::AbstractArray{T,5}, x::AbstractArray{T,5},
                     pdims::PoolDims, ::Val{K}; # == kernel_size(pdims)
-                    alpha::T=T(1), beta::T=T(0)) where {T, K}
+                    alpha::T=T(1), beta::T=T(0), kwargs...) where {T, K}
         check_dims(size(x), size(dy), pdims)
 
         width, height, depth = input_size(pdims)
@@ -182,6 +196,11 @@ for name in (:max, :mean)
             alpha = alpha / prod(K)
         end
 
+        p = if $(name != :lpnorm) 0 else
+            !haskey(kwargs, :p) && error("lpnormpool must pass p")
+            kwargs[:p]
+        end
+
         # Start with the central region
         w_region, h_region, d_region = central_region
         @inbounds for batch_idx in 1:size(x, 5), c in 1:out_c
@@ -226,6 +245,10 @@ for name in (:max, :mean)
                 elseif $(name == :mean)
                     # Either does meanpool :(
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha
+                elseif $(name == :lpnorm)
+                    # y = (∑ᵢ xᵢ^p)^(1 / p), ∂y/∂xᵢ = xᵢ^(p-1) × y^(1-p)
+                    grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
+                    dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                 else
                     error("Unimplemented codegen path")
                 end
@@ -286,6 +309,9 @@ for name in (:max, :mean)
                                 end
                             elseif $(name == :mean)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha #+ beta * dx[x_idxs...]
+                            elseif $(name == :lpnorm)
+                                grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
+                                dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                             else
                                 error("Unimplemented codegen path")
                             end

src/pooling.jl

@@ -8,11 +8,15 @@
 #     - maxpool!(y, x, pdims)
 #     - meanpool(x, pdims)
 #     - meanpool!(y, x, pdims)
+#     - lpnormpool(x, pdims)
+#     - lpnormpool!(y, x, pdims)
 #   - Pooling input backprop
 #     - ∇maxpool(dy, y, x, pdims)
 #     - ∇maxpool!(dx, dy, y, x, pdims)
 #     - ∇meanpool(dy, y, x, pdims)
 #     - ∇meanpool!(dx, dy, y, x pdims)
+#     - ∇lpnormpool(dy, y, x, pdims)
+#     - ∇lpnormpool!(dx, dy, y, x pdims)
 #
 #   All methods require a `PoolDims` object to define the dimensions and optional
 #   elements of the convolution (stride, dilation, etc...), which is easily constructable
@@ -26,6 +30,7 @@ for (front_name, backend) in (
         # This maps from public, front-facing name, to internal backend name
         :maxpool  => :direct,
         :meanpool => :direct,
+        :lpnormpool => :direct,
     )
 
     # We only define 3d pooling primitives, we reshape lower down to get 1d and 2d pooling
@@ -42,6 +47,7 @@ end
 for (front_name, backend) in (
         :∇maxpool  => :direct,
         :∇meanpool => :direct,
+        :∇lpnormpool => :direct,
     )
     @eval begin
         function $(Symbol("$(front_name)!"))(
@@ -57,7 +63,7 @@ end
 # Our strategy for pooling is to reshape to an array with three spatial dimensions, which
 # makes things MUCH EASIER for us on the backend side, and is in general pretty fast,
 # since we can specialize on sizes.
-for front_name in (:maxpool, :meanpool)
+for front_name in (:maxpool, :meanpool, :lpnormpool)
     for backend in (Symbol(), :_direct)
         for N in (3, 4)
             @eval begin
@@ -103,7 +109,7 @@ end
 # Finally, let's generate auto-allocating versions of all our functions, for all backends:
 for backend in (Symbol(), :_direct, :_nnpack)
     # First make auto-allocating versions of the basic pooling calls:
-    for name in (:maxpool, :meanpool)
+    for name in (:maxpool, :meanpool, :lpnormpool)
         @eval begin
             function $(Symbol("$(name)$(backend)"))(
                             x::AbstractArray{xT,N},
@@ -141,9 +147,15 @@ expand(N, i::Integer) = ntuple(_ -> i, N)
 
 
 """
-    maxpool(x, k::NTuple; pad=0, stride=k)
+    maxpool(x, k::NTuple{N, Integer}; pad=0, stride=k)
 
 Perform max pool operation with window size `k` on input tensor `x`.
+
+Arguments:
+
+* `x` and `k`: Expects `ndim(x) ∈ 3:5`, and always `length(k) == ndim(x) - 2`
+* `pad`: See [`pad_zeros`](@ref) for details.
+* `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`.
 """
 function maxpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
     pad = expand(Val(N), pad)
@@ -154,9 +166,15 @@ end
 
 
 """
-    meanpool(x, k::NTuple; pad=0, stride=k)
+    meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k)
 
 Perform mean pool operation with window size `k` on input tensor `x`.
+
+Arguments:
+
+* `x` and `k`: Expects `ndim(x) ∈ 3:5``, and always `length(k) == ndim(x) - 2`
+* `pad`: See [`pad_zeros`](@ref) for details.
+* `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`.
 """
 function meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
     pad = expand(Val(N), pad)
@@ -166,7 +184,33 @@ function meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
 end
 
 
-for pool in [:maxpool, :meanpool]
+"""
+    lpnormpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k)
+
+Perform Lp pool operation with value of the Lp norm `p` and window size `k` on input tensor `x`, also known as LPPool in pytorch.
+This pooling operator from [Learned-Norm Pooling for Deep Feedforward and Recurrent Neural Networks](https://arxiv.org/abs/1311.1780).
+
+Arguments:
+
+* `x` and `k`: Expects `ndim(x) ∈ 3:5``, and always `length(k) == ndim(x) - 2`
+* `p` is restricted to `0 < p < Inf`.
+* `pad`: See [`pad_zeros`](@ref) for details.
+* `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`.
+
+For all elements `x` in a size `k` window, lpnormpool computes `(∑ᵢ xᵢ^p)^(1 / p)` as an element of the output.
+
+Thus `lpnormpool(x, 1, k) ./ prod(k) ≈ meanpool(x, k)` and `lpnormpool(x, 2, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)`.
+"""
+function lpnormpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
+    (isinf(p) || p < 0) && error("p value of Lp norm pool expects `0 < p < Inf`, but p is $(p) now.")
+    pad = expand(Val(N), pad)
+    stride = expand(Val(N), stride)
+    pdims = PoolDims(x, k; padding=pad, stride=stride)
+    return lpnormpool(x, pdims; p=p)
+end
+
+
+for pool in [:maxpool, :meanpool, :lpnormpool]
     ∇pool = Symbol(:∇, pool)
     pullback = Symbol(pool, :_pullback)
     @eval function rrule(::typeof($pool), x, pdims::PoolDims; kw...)

test/perf/perf_report.jl

@@ -93,6 +93,7 @@ for rank in (2,),
     for (pool, ∇pool, name) in (
             (NNlib.maxpool!, NNlib.∇maxpool!, "maxpool"),
             (NNlib.meanpool!, NNlib.∇meanpool!, "meanpool"),
+            (NNlib.lpnormpool!, NNlib.∇lpnormpool!, "lpnormpool"),
         )
 
         t_fwd  = @benchmark $(pool)( $y, $x, $pdims)