Make all config dicts const and capitalise

Also misc. formatting and cleanup

Author: theabhirath

src/convnets/convmixer.jl

@@ -28,15 +28,15 @@ function convmixer(planes, depth; inchannels = 3, kernel_size = (9, 9),
     return Chain(Chain(stem..., Chain(blocks)), head)
 end
 
-convmixer_configs = Dict(:base => Dict(:planes => 1536, :depth => 20,
-                                       :kernel_size => (9, 9),
-                                       :patch_size => (7, 7)),
-                         :small => Dict(:planes => 768, :depth => 32,
-                                        :kernel_size => (7, 7),
-                                        :patch_size => (7, 7)),
-                         :large => Dict(:planes => 1024, :depth => 20,
-                                        :kernel_size => (9, 9),
-                                        :patch_size => (7, 7)))
+const CONVMIXER_CONFIGS = Dict(:base => Dict(:planes => 1536, :depth => 20,
+                                             :kernel_size => (9, 9),
+                                             :patch_size => (7, 7)),
+                               :small => Dict(:planes => 768, :depth => 32,
+                                              :kernel_size => (7, 7),
+                                              :patch_size => (7, 7)),
+                               :large => Dict(:planes => 1024, :depth => 20,
+                                              :kernel_size => (9, 9),
+                                              :patch_size => (7, 7)))
 
 """
     ConvMixer(mode::Symbol = :base; inchannels = 3, activation = gelu, nclasses = 1000)
@@ -57,11 +57,11 @@ end
 @functor ConvMixer
 
 function ConvMixer(mode::Symbol = :base; inchannels = 3, activation = gelu, nclasses = 1000)
-    _checkconfig(mode, keys(convmixer_configs))
-    planes = convmixer_configs[mode][:planes]
-    depth = convmixer_configs[mode][:depth]
-    kernel_size = convmixer_configs[mode][:kernel_size]
-    patch_size = convmixer_configs[mode][:patch_size]
+    _checkconfig(mode, keys(CONVMIXER_CONFIGS))
+    planes = CONVMIXER_CONFIGS[mode][:planes]
+    depth = CONVMIXER_CONFIGS[mode][:depth]
+    kernel_size = CONVMIXER_CONFIGS[mode][:kernel_size]
+    patch_size = CONVMIXER_CONFIGS[mode][:patch_size]
     layers = convmixer(planes, depth; inchannels, kernel_size, patch_size, activation,
                        nclasses)
     return ConvMixer(layers)

src/convnets/convnext.jl

@@ -66,11 +66,11 @@ function convnext(depths, planes; inchannels = 3, drop_path_rate = 0.0, λ = 1.0
 end
 
 # Configurations for ConvNeXt models
-convnext_configs = Dict(:tiny => ([3, 3, 9, 3], [96, 192, 384, 768]),
-                        :small => ([3, 3, 27, 3], [96, 192, 384, 768]),
-                        :base => ([3, 3, 27, 3], [128, 256, 512, 1024]),
-                        :large => ([3, 3, 27, 3], [192, 384, 768, 1536]),
-                        :xlarge => ([3, 3, 27, 3], [256, 512, 1024, 2048]))
+const CONVNEXT_CONFIGS = Dict(:tiny => ([3, 3, 9, 3], [96, 192, 384, 768]),
+                              :small => ([3, 3, 27, 3], [96, 192, 384, 768]),
+                              :base => ([3, 3, 27, 3], [128, 256, 512, 1024]),
+                              :large => ([3, 3, 27, 3], [192, 384, 768, 1536]),
+                              :xlarge => ([3, 3, 27, 3], [256, 512, 1024, 2048]))
 
 struct ConvNeXt
     layers::Any
@@ -94,8 +94,8 @@ See also [`Metalhead.convnext`](#).
 """
 function ConvNeXt(mode::Symbol = :base; inchannels = 3, drop_path_rate = 0.0, λ = 1.0f-6,
                   nclasses = 1000)
-    _checkconfig(mode, keys(convnext_configs))
-    layers = convnext(convnext_configs[mode]...; inchannels, drop_path_rate, λ, nclasses)
+    _checkconfig(mode, keys(CONVNEXT_CONFIGS))
+    layers = convnext(CONVNEXT_CONFIGS[mode]...; inchannels, drop_path_rate, λ, nclasses)
     return ConvNeXt(layers)
 end
 

src/convnets/densenet.jl

@@ -140,7 +140,7 @@ end
 backbone(m::DenseNet) = m.layers[1]
 classifier(m::DenseNet) = m.layers[2]
 
-const densenet_configs = Dict(121 => (6, 12, 24, 16),
+const DENSENET_CONFIGS = Dict(121 => (6, 12, 24, 16),
                               161 => (6, 12, 36, 24),
                               169 => (6, 12, 32, 32),
                               201 => (6, 12, 48, 32))
@@ -160,8 +160,8 @@ Set `pretrain = true` to load the model with pre-trained weights for ImageNet.
 See also [`Metalhead.densenet`](#).
 """
 function DenseNet(config::Integer = 121; pretrain = false, nclasses = 1000)
-    _checkconfig(config, keys(densenet_configs))
-    model = DenseNet(densenet_configs[config]; nclasses = nclasses)
+    _checkconfig(config, keys(DENSENET_CONFIGS))
+    model = DenseNet(DENSENET_CONFIGS[config]; nclasses = nclasses)
     if pretrain
         loadpretrain!(model, string("DenseNet", config))
     end

src/convnets/efficientnet.jl

@@ -59,7 +59,7 @@ end
 # e: expantion ratio
 # i: block input channels
 # o: block output channels
-const efficientnet_block_configs = [
+const EFFICIENTNET_BLOCK_CONFIGS = [
     # (n, k, s, e, i, o)
     (1, 3, 1, 1, 32, 16),
     (2, 3, 2, 6, 16, 24),
@@ -73,7 +73,7 @@ const efficientnet_block_configs = [
 # w: width scaling
 # d: depth scaling
 # r: image resolution
-const efficientnet_global_configs = Dict(:b0 => (224, (1.0, 1.0)),
+const EFFICIENTNET_GLOBAL_CONFIGS = Dict(:b0 => (224, (1.0, 1.0)),
                                          :b1 => (240, (1.0, 1.1)),
                                          :b2 => (260, (1.1, 1.2)),
                                          :b3 => (300, (1.2, 1.4)),
@@ -137,8 +137,8 @@ See also [`efficientnet`](#).
   - `pretrain`: set to `true` to load the pre-trained weights for ImageNet
 """
 function EfficientNet(name::Symbol; pretrain = false)
-    _checkconfig(name, keys(efficientnet_global_configs))
-    model = EfficientNet(efficientnet_global_configs[name][2], efficientnet_block_configs)
+    _checkconfig(name, keys(EFFICIENTNET_GLOBAL_CONFIGS))
+    model = EfficientNet(EFFICIENTNET_GLOBAL_CONFIGS[name][2], EFFICIENTNET_BLOCK_CONFIGS)
     pretrain && loadpretrain!(model, string("efficientnet-", name))
     return model
 end

src/convnets/inception/xception.jl

@@ -35,7 +35,7 @@ function xception_block(inchannels, outchannels, nrepeats; stride = 1,
         push!(layers, relu)
         append!(layers,
                 depthwise_sep_conv_norm((3, 3), inc, outc; pad = 1, bias = false,
-                                      use_bn = (false, false)))
+                                        use_norm = (false, false)))
         push!(layers, BatchNorm(outc))
     end
     layers = start_with_relu ? layers : layers[2:end]

src/convnets/mobilenet/mobilenetv1.jl

@@ -31,7 +31,7 @@ function mobilenetv1(width_mult, config;
         for _ in 1:nrepeats
             layer = dw ?
                     depthwise_sep_conv_norm((3, 3), inchannels, outch, activation;
-                                          stride = stride, pad = 1, bias = false) :
+                                            stride = stride, pad = 1, bias = false) :
                     conv_norm((3, 3), inchannels, outch, activation; stride, pad = 1,
                               bias = false)
             append!(layers, layer)
@@ -45,7 +45,7 @@ function mobilenetv1(width_mult, config;
                        Dense(inchannels, nclasses)))
 end
 
-const mobilenetv1_configs = [
+const MOBILENETV1_CONFIGS = [
     # dw, c, s, r
     (false, 32, 2, 1),
     (true, 64, 1, 1),
@@ -84,7 +84,7 @@ end
 
 function MobileNetv1(width_mult::Number = 1; inchannels = 3, pretrain = false,
                      nclasses = 1000)
-    layers = mobilenetv1(width_mult, mobilenetv1_configs; inchannels, nclasses)
+    layers = mobilenetv1(width_mult, MOBILENETV1_CONFIGS; inchannels, nclasses)
     if pretrain
         loadpretrain!(layers, string("MobileNetv1"))
     end

src/convnets/mobilenet/mobilenetv2.jl

@@ -46,7 +46,7 @@ function mobilenetv2(width_mult, configs; inchannels = 3, max_width = 1280, ncla
 end
 
 # Layer configurations for MobileNetv2
-const mobilenetv2_configs = [
+const MOBILENETV2_CONFIGS = [
     #  t,   c, n, s,     a
     (1, 16, 1, 1, relu6),
     (6, 24, 2, 2, relu6),
@@ -83,7 +83,7 @@ See also [`Metalhead.mobilenetv2`](#).
 """
 function MobileNetv2(width_mult::Number = 1; inchannels = 3, pretrain = false,
                      nclasses = 1000)
-    layers = mobilenetv2(width_mult, mobilenetv2_configs; inchannels, nclasses)
+    layers = mobilenetv2(width_mult, MOBILENETV2_CONFIGS; inchannels, nclasses)
     pretrain && loadpretrain!(layers, string("MobileNetv2"))
     if pretrain
         loadpretrain!(layers, string("MobileNetv2"))

src/convnets/mobilenet/mobilenetv3.jl

@@ -53,7 +53,7 @@ function mobilenetv3(width_mult, configs; inchannels = 3, max_width = 1024, ncla
 end
 
 # Configurations for small and large mode for MobileNetv3
-mobilenetv3_configs = Dict(:small => [
+MOBILENETV3_CONFIGS = Dict(:small => [
                                # k, t, c, SE, a, s
                                (3, 1, 16, 4, relu, 2),
                                (3, 4.5, 24, nothing, relu, 2),
@@ -115,7 +115,7 @@ function MobileNetv3(mode::Symbol = :small, width_mult::Number = 1; inchannels =
                      pretrain = false, nclasses = 1000)
     @assert mode in [:large, :small] "`mode` has to be either :large or :small"
     max_width = (mode == :large) ? 1280 : 1024
-    layers = mobilenetv3(width_mult, mobilenetv3_configs[mode]; inchannels, max_width,
+    layers = mobilenetv3(width_mult, MOBILENETV3_CONFIGS[mode]; inchannels, max_width,
                          nclasses)
     if pretrain
         loadpretrain!(layers, string("MobileNetv3", mode))

src/convnets/resnets/core.jl

@@ -25,7 +25,7 @@ function basicblock(inplanes::Integer, planes::Integer; stride::Integer = 1,
                     drop_block = identity, drop_path = identity,
                     attn_fn = planes -> identity)
     first_planes = planes ÷ reduction_factor
-    outplanes = planes * expansion_factor(basicblock)
+    outplanes = planes
     conv_bn1 = conv_norm((3, 3), inplanes => first_planes, identity; norm_layer, revnorm,
                          stride, pad = 1, bias = false)
     conv_bn2 = conv_norm((3, 3), first_planes => outplanes, identity; norm_layer, revnorm,
@@ -67,7 +67,7 @@ function bottleneck(inplanes::Integer, planes::Integer; stride::Integer,
                     attn_fn = planes -> identity)
     width = floor(Int, planes * (base_width / 64)) * cardinality
     first_planes = width ÷ reduction_factor
-    outplanes = planes * expansion_factor(bottleneck)
+    outplanes = planes * 4
     conv_bn1 = conv_norm((1, 1), inplanes => first_planes, activation; norm_layer, revnorm,
                          bias = false)
     conv_bn2 = conv_norm((3, 3), first_planes => width, identity; norm_layer, revnorm,
@@ -215,7 +215,7 @@ function basicblock_builder(block_repeats::Vector{<:Integer}; inplanes::Integer
         drop_block = DropBlock(blockschedule[schedule_idx])
         block = basicblock(inplanes, planes; stride, reduction_factor, activation,
                            norm_layer, revnorm, attn_fn, drop_path, drop_block)
-        downsample = downsample_fn(inplanes, planes * expansion; stride)
+        downsample = downsample_fn(inplanes, planes * expansion; stride, norm_layer, revnorm)
         # inplanes increases by expansion after each block
         inplanes = planes * expansion
         return block, downsample
@@ -248,18 +248,14 @@ function bottleneck_builder(block_repeats::Vector{<:Integer}; inplanes::Integer
         block = bottleneck(inplanes, planes; stride, cardinality, base_width,
                            reduction_factor, activation, norm_layer, revnorm,
                            attn_fn, drop_path, drop_block)
-        downsample = downsample_fn(inplanes, planes * expansion; stride)
+        downsample = downsample_fn(inplanes, planes * expansion; stride, norm_layer, revnorm)
         # inplanes increases by expansion after each block
         inplanes = planes * expansion
         return block, downsample
     end
     return get_layers
 end
 
-# Makes the main stages of the ResNet model. This is an internal function and should not be 
-# used by end-users. `block_fn` is a function that returns a single block of the ResNet. 
-# See `basicblock` and `bottleneck` for examples. A block must define a function 
-# `expansion(::typeof(block))` that returns the expansion factor of the block.
 function resnet_stages(get_layers, block_repeats::Vector{<:Integer}, connection)
     # Construct each stage
     stages = []
@@ -316,15 +312,15 @@ function resnet(block_type::Symbol, block_repeats::Vector{<:Integer};
     end
     classifier_fn = nfeatures -> create_classifier(nfeatures, nclasses; dropout_rate,
                                                    pool_layer, use_conv)
-    return resnet((imsize..., inchannels), stem, connection$activation, get_layers,
-                  block_repeats, classifier_fn)
+    return resnet((imsize..., inchannels), stem, get_layers, block_repeats,
+                  connection$activation, classifier_fn)
 end
 function resnet(block_fn, block_repeats, downsample_opt::Symbol = :B; kwargs...)
     return resnet(block_fn, block_repeats, shortcut_dict[downsample_opt]; kwargs...)
 end
 
 # block-layer configurations for ResNet-like models
-const resnet_configs = Dict(18 => (:basicblock, [2, 2, 2, 2]),
+const RESNET_CONFIGS = Dict(18 => (:basicblock, [2, 2, 2, 2]),
                             34 => (:basicblock, [3, 4, 6, 3]),
                             50 => (:bottleneck, [3, 4, 6, 3]),
                             101 => (:bottleneck, [3, 4, 23, 3]),

src/convnets/resnets/resnet.jl

@@ -23,8 +23,8 @@ end
 @functor ResNet
 
 function ResNet(depth::Integer; pretrain = false, inchannels = 3, nclasses = 1000)
-    _checkconfig(depth, keys(resnet_configs))
-    layers = resnet(resnet_configs[depth]...; inchannels, nclasses)
+    _checkconfig(depth, keys(RESNET_CONFIGS))
+    layers = resnet(RESNET_CONFIGS[depth]...; inchannels, nclasses)
     if pretrain
         loadpretrain!(layers, string("ResNet", depth))
     end

src/convnets/resnets/resnext.jl

@@ -29,8 +29,8 @@ end
 
 function ResNeXt(depth::Integer; pretrain = false, cardinality = 32,
                  base_width = 4, inchannels = 3, nclasses = 1000)
-    _checkconfig(depth, sort(collect(keys(resnet_configs)))[3:end])
-    layers = resnet(resnet_configs[depth]...; inchannels, nclasses, cardinality, base_width)
+    _checkconfig(depth, sort(collect(keys(RESNET_CONFIGS)))[3:end])
+    layers = resnet(RESNET_CONFIGS[depth]...; inchannels, nclasses, cardinality, base_width)
     if pretrain
         loadpretrain!(layers, string("ResNeXt", depth, "_", cardinality, "x", base_width))
     end

src/convnets/resnets/seresnet.jl

@@ -25,8 +25,8 @@ end
 (m::SEResNet)(x) = m.layers(x)
 
 function SEResNet(depth::Integer; pretrain = false, inchannels = 3, nclasses = 1000)
-    _checkconfig(depth, keys(resnet_configs))
-    layers = resnet(resnet_configs[depth]...; inchannels, nclasses,
+    _checkconfig(depth, keys(RESNET_CONFIGS))
+    layers = resnet(RESNET_CONFIGS[depth]...; inchannels, nclasses,
                     attn_fn = squeeze_excite)
     if pretrain
         loadpretrain!(layers, string("SEResNet", depth))
@@ -68,8 +68,8 @@ end
 
 function SEResNeXt(depth::Integer; pretrain = false, cardinality = 32, base_width = 4,
                    inchannels = 3, nclasses = 1000)
-    _checkconfig(depth, sort(collect(keys(resnet_configs)))[3:end])
-    layers = resnet(resnet_configs[depth]...; inchannels, nclasses, cardinality, base_width,
+    _checkconfig(depth, sort(collect(keys(RESNET_CONFIGS)))[3:end])
+    layers = resnet(RESNET_CONFIGS[depth]...; inchannels, nclasses, cardinality, base_width,
                     attn_fn = squeeze_excite)
     if pretrain
         loadpretrain!(layers, string("SEResNeXt", depth, "_", cardinality, "x", base_width))

src/convnets/vgg.jl

@@ -99,12 +99,12 @@ function vgg(imsize; config, inchannels, batchnorm = false, nclasses, fcsize, dr
     return Chain(Chain(conv), class)
 end
 
-const vgg_conv_configs = Dict(:A => [(64, 1), (128, 1), (256, 2), (512, 2), (512, 2)],
+const VGG_CONV_CONFIGS = Dict(:A => [(64, 1), (128, 1), (256, 2), (512, 2), (512, 2)],
                               :B => [(64, 2), (128, 2), (256, 2), (512, 2), (512, 2)],
                               :D => [(64, 2), (128, 2), (256, 3), (512, 3), (512, 3)],
                               :E => [(64, 2), (128, 2), (256, 4), (512, 4), (512, 4)])
 
-const vgg_configs = Dict(11 => :A,
+const VGG_CONFIGS = Dict(11 => :A,
                          13 => :B,
                          16 => :D,
                          19 => :E)
@@ -153,8 +153,8 @@ See also [`VGG`](#).
   - `pretrain`: set to `true` to load pre-trained model weights for ImageNet
 """
 function VGG(depth::Integer = 16; pretrain = false, batchnorm = false, nclasses = 1000)
-    _checkconfig(depth, keys(vgg_configs))
-    model = VGG((224, 224); config = vgg_conv_configs[vgg_configs[depth]],
+    _checkconfig(depth, keys(VGG_CONFIGS))
+    model = VGG((224, 224); config = VGG_CONV_CONFIGS[VGG_CONFIGS[depth]],
                 inchannels = 3,
                 batchnorm = batchnorm,
                 nclasses = nclasses,