Add DropBlock

Author: theabhirath

Project.toml

@@ -5,12 +5,14 @@ version = "0.7.3-DEV"
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]

src/convnets/densenet.jl

@@ -100,7 +100,7 @@ Create a DenseNet model
   - `reduction`: the factor by which the number of feature maps is scaled across each transition
   - `nclasses`: the number of output classes
 """
-function densenet(nblocks; growth_rate = 32, reduction = 0.5, nclasses = 1000)
+function densenet(nblocks::NTuple{N, <:Integer}; growth_rate = 32, reduction = 0.5, nclasses = 1000)
     return densenet(2 * growth_rate, [fill(growth_rate, n) for n in nblocks];
                     reduction = reduction, nclasses = nclasses)
 end

src/layers/Layers.jl

@@ -1,19 +1,21 @@
 module Layers
 
 using Flux
-using Flux: outputsize, Zygote
+using NNlib
+using NNlibCUDA
 using Functors
+using ChainRulesCore
 using Statistics
 using MLUtils
 
 include("../utilities.jl")
 
 include("attention.jl")
 include("embeddings.jl")
-include("mlp.jl")
+include("mlp-linear.jl")
 include("normalise.jl")
 include("conv.jl")
-include("others.jl")
+include("drop.jl")
 
 export MHAttention,
        PatchEmbedding, ViPosEmbedding, ClassTokens,

src/layers/drop.jl

@@ -0,0 +1,61 @@
+"""
+    DropBlock(drop_prob = 0.1, block_size = 7)
+
+Implements DropBlock, a regularization method for convolutional networks.
+([reference](https://arxiv.org/pdf/1810.12890.pdf))
+"""
+struct DropBlock{F}
+    drop_prob::F
+    block_size::Integer
+end
+@functor DropBlock
+
+(m::DropBlock)(x) = dropblock(x, m.drop_prob, m.block_size)
+
+DropBlock(drop_prob = 0.1, block_size = 7) = DropBlock(drop_prob, block_size)
+
+function _dropblock_checks(x, drop_prob, T)
+    if !(T <: AbstractArray)
+        throw(ArgumentError("x must be an `AbstractArray`"))
+    end
+    if ndims(x) != 4
+        throw(ArgumentError("x must have 4 dimensions (H, W, C, N) for `DropBlock`"))
+    end
+    @assert drop_prob < 0 || drop_prob > 1 "drop_prob must be between 0 and 1, got $drop_prob"
+end
+ChainRulesCore.@non_differentiable _dropblock_checks(x, drop_prob, T)
+
+function dropblock(x::T, drop_prob, block_size::Integer) where {T}
+    _dropblock_checks(x, drop_prob, T)
+    if drop_prob == 0
+        return x
+    end
+    return _dropblock(x, drop_prob, block_size)
+end
+
+function _dropblock(x::AbstractArray{T, 4}, drop_prob, block_size) where {T}
+    gamma = drop_prob / (block_size ^ 2)
+    mask = rand_like(x, Float32, (size(x, 1), size(x, 2), size(x, 3)))
+    mask .<= gamma
+    block_mask = maxpool(reshape(mask, (size(mask)[1:3]..., 1)), (block_size, block_size);
+                         pad = block_size ÷ 2, stride = (1, 1))
+    if block_size % 2 == 0
+        block_mask = block_mask[1:(end - 1), 1:(end - 1), :, :]
+    end
+    block_mask = 1 .- dropdims(block_mask; dims = 4)
+    out = (x .* reshape(block_mask, (size(block_mask)[1:3]..., 1))) * length(block_mask) /
+          sum(block_mask)
+    return out
+end
+
+"""
+    DropPath(p)
+
+Implements Stochastic Depth - equivalent to `Dropout(p; dims = 4)` when `p` ≥ 0.
+([reference](https://arxiv.org/abs/1603.09382))
+
+# Arguments
+
+  - `p`: rate of Stochastic Depth.
+"""
+DropPath(p) = p ≥ 0 ? Dropout(p; dims = 4) : identity

src/layers/mlp.jl renamed to src/layers/mlp-linear.jl

@@ -1,3 +1,18 @@
+"""
+    LayerScale(λ, planes::Integer)
+
+Creates a `Flux.Scale` layer that performs "`LayerScale`"
+([reference](https://arxiv.org/abs/2103.17239)).
+
+# Arguments
+
+  - `planes`: Size of channel dimension in the input.
+  - `λ`: initialisation value for the learnable diagonal matrix.
+"""
+function LayerScale(planes::Integer, λ)
+    return λ > 0 ? Flux.Scale(fill(Float32(λ), planes), false) : identity
+end
+
 """
     mlp_block(inplanes::Integer, hidden_planes::Integer, outplanes::Integer = inplanes; 
               dropout = 0., activation = gelu)