Add ResNeXt back

Also add tests. A lot of tests

Author: theabhirath

src/Metalhead.jl

@@ -38,7 +38,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,
@@ -47,7 +47,7 @@ export AlexNet, VGG, VGG11, VGG13, VGG16, VGG19,
        ConvMixer, ConvNeXt
 
 # use Flux._big_show to pretty print large models
-for T in (:AlexNet, :VGG, :DenseNet, :ResNet, # :ResNeXt, 
+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

@@ -173,18 +173,18 @@ function _drop_blocks(drop_block_prob = 0.0)
     ]
 end
 
-function resnet(block, layers; nclasses = 1000, inchannels = 3, output_stride = 32,
+function resnet(block_fn, layers; nclasses = 1000, inchannels = 3, output_stride = 32,
                 stem = first(resnet_stem(; inchannels)), inplanes = 64,
                 downsample_fn = downsample_conv,
                 drop_rates::NamedTuple = (drop_rate = 0.0, drop_path_rate = 0.0,
                                           drop_block_rate = 0.0),
                 block_args::NamedTuple = NamedTuple())
     # Feature Blocks
     channels = [64, 128, 256, 512]
-    stage_blocks = _make_blocks(block, channels, layers, inplanes;
+    stage_blocks = _make_blocks(block_fn, channels, layers, inplanes;
                                 output_stride, downsample_fn, drop_rates, block_args)
     # Head (Pooling and Classifier)
-    expansion = expansion_factor(block)
+    expansion = expansion_factor(block_fn)
     num_features = 512 * expansion
     classifier = Chain(GlobalMeanPool(), Dropout(drop_rates.drop_rate), MLUtils.flatten,
                        Dense(num_features, nclasses))
@@ -201,11 +201,27 @@ struct ResNet
 end
 @functor ResNet
 
-function ResNet(depth::Integer; pretrain = false, nclasses = 1000, kwargs...)
-    @assert depth in [18, 34, 50, 101, 152] "Invalid depth. Must be one of [18, 34, 50, 101, 152]"
-    model = resnet(resnet_config[depth]...; nclasses, kwargs...)
+function ResNet(depth::Integer; pretrain = false, nclasses = 1000)
+    @assert depth in [18, 34, 50, 101, 152]
+    "Invalid depth. Must be one of [18, 34, 50, 101, 152]"
+    model = resnet(resnet_config[depth]...; nclasses)
     if pretrain
         loadpretrain!(model, string("resnet", depth))
     end
     return model
 end
+
+struct ResNeXt
+    layers::Any
+end
+@functor ResNeXt
+
+function ResNeXt(depth::Integer; pretrain = false, cardinality = 32, base_width = 4, nclasses = 1000)
+    @assert depth in [50, 101, 152]
+    "Invalid depth. Must be one of [50, 101, 152]"
+    model = resnet(bottleneck, [3, 4, 6, 3]; nclasses, block_args = (; cardinality, base_width))
+    if pretrain
+        loadpretrain!(model, string("resnext", depth, "_", cardinality, "x", base_width))
+    end
+    return model
+end

src/layers/Layers.jl

@@ -2,8 +2,7 @@ module Layers
 
 using Flux
 using CUDA
-using NNlib
-using NNlibCUDA
+using NNlib, NNlibCUDA
 using Functors
 using ChainRulesCore
 using Statistics

src/utilities.jl

@@ -9,24 +9,6 @@ function _round_channels(channels, divisor, min_value = divisor)
     return (new_channels < 0.9 * channels) ? new_channels + divisor : new_channels
 end
 
-"""
-    addrelu(x, y)
-
-Convenience function for `(x, y) -> @. relu(x + y)`.
-Useful as the `connection` argument for [`resnet`](#).
-See also [`reluadd`](#).
-"""
-addrelu(x, y) = @. relu(x + y)
-
-"""
-    reluadd(x, y)
-
-Convenience function for `(x, y) -> @. relu(x) + relu(y)`.
-Useful as the `connection` argument for [`resnet`](#).
-See also [`addrelu`](#).
-"""
-reluadd(x, y) = @. relu(x) + relu(y)
-
 """
     cat_channels(x, y, zs...)
 

test/convnets.jl

@@ -27,18 +27,39 @@ GC.safepoint()
 GC.gc()
 
 @testset "ResNet" begin
-    @testset "ResNet($sz)" for sz in [18, 34, 50, 101, 152]
-        m = ResNet(sz)
-        @test size(m(x_256)) == (1000, 1)
-        ## TODO: find a way to port pretrained models to the new ResNet API
+    # Tests for pretrained ResNets
+    ## TODO: find a way to port pretrained models to the new ResNet API
+    # @testset "ResNet($sz)" for sz in [18, 34, 50, 101, 152]
         # if (ResNet, sz) in PRETRAINED_MODELS
         #     @test acctest(ResNet(sz, pretrain = true))
         # else
         #     @test_throws ArgumentError ResNet(sz, pretrain = true)
         # end
-        @test gradtest(m, x_256)
-        GC.safepoint()
-        GC.gc()
+    # end
+
+    @testset "resnet" begin
+        @testset for block_fn in [Metalhead.basicblock, Metalhead.bottleneck]
+            layer_list = [
+                [2, 2, 2, 2],
+                [3, 4, 6, 3],
+                [3, 4, 23, 3],
+                [3, 8, 36, 3]
+            ]
+            @testset for layers in layer_list
+                drop_list = [
+                    (drop_rate = 0.1, drop_path_rate = 0.1, drop_block_rate = 0.1),
+                    (drop_rate = 0.5, drop_path_rate = 0.5, drop_block_rate = 0.5),
+                    (drop_rate = 0.9, drop_path_rate = 0.9, drop_block_rate = 0.9),
+                ]
+                @testset for drop_rates in drop_list
+                    m = Metalhead.resnet(block_fn, layers; drop_rates)
+                    @test size(m(x_224)) == (1000, 1)
+                    @test gradtest(m, x_224)
+                    GC.safepoint()
+                    GC.gc()
+                end
+            end
+        end
     end
 end
 
@@ -47,16 +68,20 @@ GC.gc()
 
 @testset "ResNeXt" begin
     @testset for depth in [50, 101, 152]
-        m = ResNeXt(depth)
-        @test size(m(x_224)) == (1000, 1)
-        if ResNeXt in PRETRAINED_MODELS
-            @test acctest(ResNeXt(depth, pretrain = true))
-        else
-            @test_throws ArgumentError ResNeXt(depth, pretrain = true)
+        @testset for cardinality in [32, 64]
+            @testset for base_width in [4, 8]
+                m = ResNeXt(depth; cardinality, base_width)
+                @test size(m(x_224)) == (1000, 1)
+                if string("resnext", depth, "_", cardinality, "x", base_width) in PRETRAINED_MODELS
+                    @test acctest(ResNeXt(depth, pretrain = true))
+                else
+                    @test_throws ArgumentError ResNeXt(depth, pretrain = true)
+                end
+                @test gradtest(m, x_224)
+                GC.safepoint()
+                GC.gc()
+            end
         end
-        @test gradtest(m, x_224)
-        GC.safepoint()
-        GC.gc()
     end
 end