Add WideResNet

Author: theabhirath

.github/workflows/CI.yml

@@ -29,7 +29,8 @@ jobs:
           - '["AlexNet", "VGG"]'
           - '["GoogLeNet", "SqueezeNet", "MobileNet"]'
           - '["EfficientNet"]'
-          - '[r"ResNet", r"ResNeXt"]'
+          - 'r"/^ResNet\z/"'
+          - '[r"ResNeXt", r"SEResNet"]'
           - '"Inception"'
           - '"DenseNet"'
           - '["ConvNeXt", "ConvMixer"]'

src/Metalhead.jl

@@ -60,7 +60,7 @@ export AlexNet, VGG, VGG11, VGG13, VGG16, VGG19,
        DenseNet, DenseNet121, DenseNet161, DenseNet169, DenseNet201,
        GoogLeNet, Inception3, Inceptionv3, Inceptionv4, InceptionResNetv2, Xception,
        SqueezeNet, MobileNetv1, MobileNetv2, MobileNetv3, EfficientNet,
-       SEResNet, SEResNeXt,
+       WideResNet, SEResNet, SEResNeXt,
        MLPMixer, ResMLP, gMLP,
        ViT,
        ConvMixer, ConvNeXt

src/convnets/resnets/resnet.jl

@@ -35,3 +35,43 @@ end
 
 backbone(m::ResNet) = m.layers[1]
 classifier(m::ResNet) = m.layers[2]
+
+"""
+    WideResNet(depth::Integer; pretrain = false, inchannels = 3, nclasses = 1000)
+
+Creates a Wide ResNet model with the specified depth. The model is the same as ResNet
+except for the bottleneck number of channels which is twice larger in every block.
+The number of channels in outer 1x1 convolutions is the same.
+((reference)[https://arxiv.org/abs/1605.07146])
+
+# Arguments
+
+  - `depth`: one of `[18, 34, 50, 101, 152]`. The depth of the Wide ResNet model.
+  - `pretrain`: set to `true` to load the model with pre-trained weights for ImageNet
+  - `inchannels`: The number of input channels.
+  - `nclasses`: the number of output classes
+
+!!! warning
+    
+    `WideResNet` does not currently support pretrained weights.
+
+Advanced users who want more configuration options will be better served by using [`resnet`](#).
+"""
+struct WideResNet
+    layers::Any
+end
+@functor WideResNet
+
+function WideResNet(depth::Integer; pretrain = false, inchannels = 3, nclasses = 1000)
+    _checkconfig(depth, [50, 101])
+    layers = resnet(RESNET_CONFIGS[depth]...; base_width = 128, inchannels, nclasses)
+    if pretrain
+        loadpretrain!(layers, string("WideResNet", depth))
+    end
+    return WideResNet(layers)
+end
+
+(m::WideResNet)(x) = m.layers(x)
+
+backbone(m::WideResNet) = m.layers[1]
+classifier(m::WideResNet) = m.layers[2]

src/layers/conv.jl

@@ -52,7 +52,8 @@ function conv_norm(kernel_size, inplanes::Int, outplanes::Int, activation = relu
     return revnorm ? reverse(layers) : layers
 end
 
-function conv_norm(kernel_size, ch::Pair{<:Integer, <:Integer}, activation = identity; kwargs...)
+function conv_norm(kernel_size, ch::Pair{<:Integer, <:Integer}, activation = identity;
+                   kwargs...)
     inplanes, outplanes = ch
     return conv_norm(kernel_size, inplanes, outplanes, activation; kwargs...)
 end

src/vit-based/vit.jl

@@ -16,7 +16,8 @@ function transformer_encoder(planes, depth, nheads; mlp_ratio = 4.0, dropout_rat
     layers = [Chain(SkipConnection(prenorm(planes,
                                            MHAttention(planes, nheads;
                                                        attn_dropout_rate = dropout_rate,
-                                                       proj_dropout_rate = dropout_rate)), +),
+                                                       proj_dropout_rate = dropout_rate)),
+                                   +),
                     SkipConnection(prenorm(planes,
                                            mlp_block(planes, floor(Int, mlp_ratio * planes);
                                                      dropout_rate)), +))

test/convnets.jl

@@ -56,6 +56,15 @@ end
             end
         end
     end
+
+    @testset "WideResNet" begin
+        @testset "WideResNet($sz)" for sz in [50, 101]
+            m = WideResNet(sz)
+            @test size(m(x_224)) == (1000, 1)
+            @test gradtest(m, x_224)
+            _gc()
+        end
+    end
 end
 
 @testset "ResNeXt" begin