diff --git a/NEWS.md b/NEWS.md index 587d909f87..46f410eb0c 100644 --- a/NEWS.md +++ b/NEWS.md @@ -2,6 +2,7 @@ * Change to `DataLoader`'s constructor [https://github.com/FluxML/Flux.jl/pull/1152] * Use `DataLoader` with `NamedTuple`s, so that tensors can be accessed by name [https://github.com/FluxML/Flux.jl/pull/1221]. * Error if Dense layers weights and biases are not arrays [https://github.com/FluxML/Flux.jl/pull/1218]. +* Add `Adaptive Pooling` in Flux layers [https://github.com/FluxML/Flux.jl/pull/1239]. # v0.10.5 * Add option for [same padding](https://github.com/FluxML/Flux.jl/pull/901) to conv and pooling layers by setting `pad=SamePad()`. diff --git a/docs/src/models/layers.md b/docs/src/models/layers.md index d228cad8bd..87178536ae 100644 --- a/docs/src/models/layers.md +++ b/docs/src/models/layers.md @@ -13,8 +13,10 @@ These layers are used to build convolutional neural networks (CNNs). ```@docs Conv +AdaptiveMaxPool MaxPool GlobalMaxPool +AdaptiveMeanPool MeanPool GlobalMeanPool DepthwiseConv diff --git a/src/Flux.jl b/src/Flux.jl index c28a7d3637..e40f9ea7ce 100644 --- a/src/Flux.jl +++ b/src/Flux.jl @@ -12,9 +12,10 @@ using Zygote: Params, @adjoint, gradient, pullback, @nograd export gradient export Chain, Dense, Maxout, RNN, LSTM, GRU, SamePad, Conv, CrossCor, ConvTranspose, - GlobalMaxPool, GlobalMeanPool, MaxPool, MeanPool, flatten, - DepthwiseConv, Dropout, AlphaDropout, LayerNorm, BatchNorm, InstanceNorm, GroupNorm, - SkipConnection, params, fmap, cpu, gpu, f32, f64, testmode!, trainmode! + AdaptiveMaxPool, AdaptiveMeanPool, GlobalMaxPool, GlobalMeanPool, MaxPool, + MeanPool, flatten, DepthwiseConv, Dropout, AlphaDropout, LayerNorm, BatchNorm, + InstanceNorm, GroupNorm, SkipConnection, params, fmap, cpu, gpu, f32, f64, + testmode!, trainmode! include("optimise/Optimise.jl") using .Optimise diff --git a/src/layers/conv.jl b/src/layers/conv.jl index 3410ce653a..3a5bb990fb 100644 --- a/src/layers/conv.jl +++ b/src/layers/conv.jl @@ -480,6 +480,54 @@ end outdims(l::CrossCor, isize) = output_size(DenseConvDims(_paddims(isize, size(l.weight)), size(l.weight); stride = l.stride, padding = l.pad, dilation = l.dilation)) +""" + AdaptiveMaxPool(out) + +Adaptive max pooling layer. `out` is the desired output size (batch and channel dimension excluded). +""" +struct AdaptiveMaxPool{S, O} + out::NTuple{O, Int} + AdaptiveMaxPool(out::NTuple{O, Int}) where O = new{O + 2, O}(out) +end + +function (a::AdaptiveMaxPool{S})(x::AbstractArray{T, S}) where {S, T} + insize = size(x)[1:end-2] + outsize = a.out + stride = insize .÷ outsize + k = insize .- (outsize .- 1) .* stride + pad = 0 + pdims = PoolDims(x, k; padding=pad, stride=stride) + return maxpool(x, pdims) +end + +function Base.show(io::IO, a::AdaptiveMaxPool) + print(io, "AdaptiveMaxPool(", a.out, ")") +end + +""" + AdaptiveMeanPool(out) + +Adaptive mean pooling layer. `out` is the desired output size (batch and channel dimension excluded). +""" +struct AdaptiveMeanPool{S, O} + out::NTuple{O, Int} + AdaptiveMeanPool(out::NTuple{O, Int}) where O = new{O + 2, O}(out) +end + +function (a::AdaptiveMeanPool{S})(x::AbstractArray{T, S}) where {S, T} + insize = size(x)[1:end-2] + outsize = a.out + stride = insize .÷ outsize + k = insize .- (outsize .- 1) .* stride + pad = 0 + pdims = PoolDims(x, k; padding=pad, stride=stride) + return meanpool(x, pdims) +end + +function Base.show(io::IO, a::AdaptiveMeanPool) + print(io, "AdaptiveMeanPool(", a.out, ")") +end + """ GlobalMaxPool() diff --git a/test/cuda/layers.jl b/test/cuda/layers.jl index e48a464280..b51320ff4c 100644 --- a/test/cuda/layers.jl +++ b/test/cuda/layers.jl @@ -52,6 +52,9 @@ gradtest("Conv", conv_layers, r, (2,2), 1=>3) pooling_layers = [MaxPool, MeanPool] gradtest("Pooling", pooling_layers, r, (2,2)) +adaptive_pooling_layers = [AdaptiveMaxPool, AdaptiveMeanPool] +gradtest("AdaptivePooling", adaptive_pooling_layers, r, (7,7)) + dropout_layers = [Dropout, AlphaDropout] gradtest("Dropout", dropout_layers, r, 0.5f0) diff --git a/test/layers/conv.jl b/test/layers/conv.jl index 8c825bfda2..4a468c84d3 100644 --- a/test/layers/conv.jl +++ b/test/layers/conv.jl @@ -4,6 +4,15 @@ using Flux: gradient @testset "Pooling" begin x = randn(Float32, 10, 10, 3, 2) + y = randn(Float32, 20, 20, 3, 2) + ampx = AdaptiveMaxPool((5,5)) + @test ampx(x) == maxpool(x, PoolDims(x, 2)) + ampx = AdaptiveMeanPool((5,5)) + @test ampx(x) == meanpool(x, PoolDims(x, 2)) + ampy = AdaptiveMaxPool((10, 5)) + @test ampy(y) == maxpool(y, PoolDims(y, (2, 4))) + ampy = AdaptiveMeanPool((10, 5)) + @test ampy(y) == meanpool(y, PoolDims(y, (2, 4))) gmp = GlobalMaxPool() @test size(gmp(x)) == (1, 1, 3, 2) gmp = GlobalMeanPool()