Hardcode large ResNet model dict for block configs

Also misc. cleanup

Co-Authored-By: Kyle Daruwalla <[email protected]>

Author: theabhirath

src/convnets/resnets/core.jl

@@ -116,7 +116,7 @@ function downsample_identity(inplanes::Integer, outplanes::Integer; kwargs...)
 end
 
 # Shortcut configurations for the ResNet models
-const SHORTCUT_DICT = Dict(:A => (downsample_identity, downsample_identity),
+const RESNET_SHORTCUTS = Dict(:A => (downsample_identity, downsample_identity),
                            :B => (downsample_conv, downsample_identity),
                            :C => (downsample_conv, downsample_conv),
                            :D => (downsample_pool, downsample_identity))
@@ -343,7 +343,7 @@ function resnet(block_type::Symbol, block_repeats::AbstractVector{<:Integer},
                   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...)
+    return resnet(block_fn, block_repeats, RESNET_SHORTCUTS[downsample_opt]; kwargs...)
 end
 
 # block-layer configurations for ResNet-like models
@@ -352,3 +352,7 @@ const RESNET_CONFIGS = Dict(18 => (:basicblock, [2, 2, 2, 2]),
                             50 => (:bottleneck, [3, 4, 6, 3]),
                             101 => (:bottleneck, [3, 4, 23, 3]),
                             152 => (:bottleneck, [3, 8, 36, 3]))
+
+const LRESNET_CONFIGS = Dict(50 => (:bottleneck, [3, 4, 6, 3]),
+                             101 => (:bottleneck, [3, 4, 23, 3]),
+                             152 => (:bottleneck, [3, 8, 36, 3]))

src/convnets/resnets/res2net.jl

@@ -33,7 +33,11 @@ function bottle2neck(inplanes::Integer, planes::Integer; stride::Integer = 1,
                 for _ in 1:max(1, scale - 1)]
     reslayer = is_first ? Parallel(cat_channels, pool, conv_bns...) :
                Parallel(cat_channels, identity, Chain(PairwiseFusion(+, conv_bns...)))
-    tuplify(x) = is_first ? tuple(x...) : tuple(x[1], tuple(x[2:end]...))
+    if is_first
+        tuplify(x) = tuple(x...)
+    else
+        tuplify(x) = tuple(x[1], tuple(x[2:end]...))
+    end
     layers = [conv_norm((1, 1), inplanes => width * scale, activation;
                         norm_layer, revnorm, bias = false)...,
               chunk$(; size = width, dims = 3), tuplify, reslayer,
@@ -138,8 +142,8 @@ end
 function Res2NeXt(depth::Integer; pretrain::Bool = false, scale::Integer = 4,
                   base_width::Integer = 4, cardinality::Integer = 8,
                   inchannels::Integer = 3, nclasses::Integer = 1000)
-    _checkconfig(depth, sort(collect(keys(RESNET_CONFIGS)))[3:end])
-    layers = resnet(:bottle2neck, RESNET_CONFIGS[depth][2]; base_width, scale,
+    _checkconfig(depth, keys(LRESNET_CONFIGS))
+    layers = resnet(:bottle2neck, LRESNET_CONFIGS[depth][2]; base_width, scale,
                     cardinality, inchannels, nclasses)
     if pretrain
         loadpretrain!(layers,
@@ -152,4 +156,4 @@ end
 (m::Res2NeXt)(x) = m.layers(x)
 
 backbone(m::Res2NeXt) = m.layers[1]
-classifier(m::Res2NeXt) = m.layers[2]
+classifier(m::Res2NeXt) = m.layers[2]

src/convnets/resnets/resnet.jl

@@ -65,8 +65,8 @@ end
 
 function WideResNet(depth::Integer; pretrain::Bool = false, inchannels::Integer = 3,
                     nclasses::Integer = 1000)
-    _checkconfig(depth, sort(collect(keys(RESNET_CONFIGS)))[3:end])
-    layers = resnet(RESNET_CONFIGS[depth]...; base_width = 128, inchannels, nclasses)
+    _checkconfig(depth, keys(LRESNET_CONFIGS))
+    layers = resnet(LRESNET_CONFIGS[depth]...; base_width = 128, inchannels, nclasses)
     if pretrain
         loadpretrain!(layers, string("WideResNet", depth))
     end

src/convnets/resnets/resnext.jl

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

src/convnets/resnets/seresnet.jl

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