Make DropBlock really work

Author: theabhirath

Project.toml

@@ -5,14 +5,15 @@ version = "0.7.3"
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 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"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]

src/Metalhead.jl

@@ -39,7 +39,7 @@ include("vit-based/vit.jl")
 include("pretrain.jl")
 
 export AlexNet, VGG, VGG11, VGG13, VGG16, VGG19,
-       ResNet, ResNet18, ResNet34, ResNet50, ResNet101, ResNet152, ResNeXt,
+       ResNet, ResNet18, ResNet34, ResNet50, ResNet101, ResNet152, # ResNeXt,
        DenseNet, DenseNet121, DenseNet161, DenseNet169, DenseNet201,
        GoogLeNet, Inception3, Inceptionv3, Inceptionv4, InceptionResNetv2, Xception,
        SqueezeNet, MobileNetv1, MobileNetv2, MobileNetv3, EfficientNet,
@@ -48,7 +48,7 @@ export AlexNet, VGG, VGG11, VGG13, VGG16, VGG19,
        ConvMixer, ConvNeXt
 
 # use Flux._big_show to pretty print large models
-for T in (:AlexNet, :VGG, :ResNeXt, :DenseNet, :ResNet,
+for T in (:AlexNet, :VGG, :DenseNet, :ResNet, # :ResNeXt, 
           :GoogLeNet, :Inceptionv3, :Inceptionv4, :InceptionResNetv2, :Xception,
           :SqueezeNet, :MobileNetv1, :MobileNetv2, :MobileNetv3,
           :MLPMixer, :ResMLP, :gMLP, :ViT, :ConvMixer, :ConvNeXt)

src/convnets/resne(x)t.jl

@@ -107,7 +107,7 @@ function resnet_stem(; stem_type = :default, inchannels = 3, replace_stem_pool =
     else
         stempool = MaxPool((3, 3); stride = 2, pad = 1)
     end
-    return Chain(conv1, bn1, stempool)
+    return inplanes, Chain(conv1, bn1, stempool)
 end
 
 function downsample_block(downsample_fn, inplanes, planes, expansion; kernel_size = (1, 1),
@@ -150,7 +150,7 @@ function _make_blocks(block_fn, channels, block_repeats, inplanes; output_stride
         end
         # Downsample block; either a (default) convolution-based block or a pooling-based block.
         downsample = downsample_block(downsample_fn, inplanes, planes, expansion;
-                                      downsample_args...)
+                                      stride, dilation, first_dilation = dilation, downsample_args...)
         # Construct the blocks for each stage
         blocks = []
         for block_idx in 1:num_blocks
@@ -172,16 +172,16 @@ function _make_blocks(block_fn, channels, block_repeats, inplanes; output_stride
 end
 
 function resnet(block, layers; nclasses = 1000, inchannels = 3, output_stride = 32,
-                stem_fn = resnet_stem, stem_args::NamedTuple = (),
-                downsample_fn = downsample_conv, downsample_args::NamedTuple = (),
+                stem_fn = resnet_stem, stem_args::NamedTuple = NamedTuple(),
+                downsample_fn = downsample_conv, downsample_args::NamedTuple = NamedTuple(),
                 drop_rates::NamedTuple = (drop_rate = 0.0, drop_path_rate = 0.0,
-                                          drop_block_rate = 0.0),
-                block_args::NamedTuple = ())
+                                          drop_block_rate = 0.5),
+                block_args::NamedTuple = NamedTuple())
     # Stem
-    stem = stem_fn(; inchannels, stem_args...)
+    inplanes, stem = stem_fn(; inchannels, stem_args...)
     # Feature Blocks
     channels = [64, 128, 256, 512]
-    stage_blocks = _make_blocks(block, channels, layers, inchannels;
+    stage_blocks = _make_blocks(block, channels, layers, inplanes;
                                 output_stride, downsample_fn, downsample_args,
                                 drop_block_rate = drop_rates.drop_block_rate,
                                 drop_path_rate = drop_rates.drop_path_rate,

src/layers/Layers.jl

@@ -1,12 +1,14 @@
 module Layers
 
 using Flux
+using CUDA
 using NNlib
 using NNlibCUDA
 using Functors
 using ChainRulesCore
 using Statistics
 using MLUtils
+using Random
 
 include("../utilities.jl")
 

src/layers/drop.jl

@@ -1,28 +1,33 @@
-"""
-    DropBlock(drop_block_prob = 0.1, block_size = 7, gamma_scale = 1.0)
+function dropblock(rng::AbstractRNG, x::AbstractArray{T, 4}, drop_block_prob, block_size,
+                   gamma_scale, active::Bool = true) where {T}
+    active || return x
+    H, W, _, _ = size(x)
+    total_size = H * W
+    clipped_block_size = min(block_size, min(H, W))
+    gamma = gamma_scale * drop_block_prob * total_size / clipped_block_size^2 /
+            ((W - block_size + 1) * (H - block_size + 1))
+    block_mask = rand_like(rng, x) .< gamma
+    block_mask = maxpool(block_mask, (clipped_block_size, clipped_block_size);
+                         stride = 1, pad = clipped_block_size ÷ 2)
+    block_mask = 1 .- block_mask
+    normalize_scale = convert(T, (length(block_mask) / sum(block_mask) .+ 1e-6))
+    return x .* block_mask .* normalize_scale
+end
+dropout_mask(rng::CUDA.RNG, x::CuArray, p; kwargs...) = _dropout_mask(rng, x, p; kwargs...)
+function dropblock(rng, x::CuArray, p; kwargs...)
+    throw(ArgumentError("x isa CuArray, but rng isa $(typeof(rng)). dropblock only support CUDA.RNG for CuArrays."))
+end
 
-Implements DropBlock, a regularization method for convolutional networks.
-([reference](https://arxiv.org/pdf/1810.12890.pdf))
-"""
-struct DropBlock{F}
+struct DropBlock{F, R <: AbstractRNG}
     drop_block_prob::F
     block_size::Integer
     gamma_scale::F
+    active::Union{Bool, Nothing}
+    rng::R
 end
-@functor DropBlock
-
-(m::DropBlock)(x) = dropblock(x, m.drop_block_prob, m.block_size, m.gamma_scale)
 
-function DropBlock(drop_block_prob = 0.1, block_size = 7, gamma_scale = 1.0)
-    if drop_block_prob == 0.0
-        return identity
-    end
-    @assert drop_block_prob < 0 || drop_block_prob > 1
-    "drop_block_prob must be between 0 and 1, got $drop_block_prob"
-    @assert gamma_scale < 0 || gamma_scale > 1
-    "gamma_scale must be between 0 and 1, got $gamma_scale"
-    return DropBlock(drop_block_prob, block_size, gamma_scale)
-end
+@functor DropBlock
+trainable(a::DropBlock) = (;)
 
 function _dropblock_checks(x::T) where {T}
     if !(T <: AbstractArray)
@@ -34,20 +39,26 @@ function _dropblock_checks(x::T) where {T}
 end
 ChainRulesCore.@non_differentiable _dropblock_checks(x)
 
-function dropblock(x::AbstractArray{T, 4}, drop_block_prob, block_size,
-                   gamma_scale) where {T}
+function (m::DropBlock)(x)
     _dropblock_checks(x)
-    H, W, _, _ = size(x)
-    total_size = H * W
-    clipped_block_size = min(block_size, min(H, W))
-    gamma = gamma_scale * drop_block_prob * total_size / clipped_block_size^2 /
-            ((W - block_size + 1) * (H - block_size + 1))
-    block_mask = rand_like(x) .< gamma
-    block_mask = maxpool(block_mask, (clipped_block_size, clipped_block_size);
-                         stride = 1, pad = clipped_block_size ÷ 2)
-    block_mask = 1 .- block_mask
-    normalize_scale = convert(T, (length(block_mask) / sum(block_mask) .+ 1e-6))
-    return x .* block_mask .* normalize_scale
+    Flux._isactive(m) || return x
+    return dropblock(m.rng, x, m.drop_block_prob, m.block_size, m.gamma_scale)
+end
+
+function Flux.testmode!(m::DropBlock, mode = true)
+    return (m.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
+end
+
+function DropBlock(drop_block_prob = 0.1, block_size = 7, gamma_scale = 1.0,
+                   rng = Flux.rng_from_array())
+    if drop_block_prob == 0.0
+        return identity
+    end
+    @assert 0 ≤ drop_block_prob ≤ 1
+    "drop_block_prob must be between 0 and 1, got $drop_block_prob"
+    @assert 0 ≤ gamma_scale ≤ 1
+    "gamma_scale must be between 0 and 1, got $gamma_scale"
+    return DropBlock(drop_block_prob, block_size, gamma_scale, nothing, rng)
 end
 
 """