Fixup and tweaks

Author: theabhirath

src/convnets/inception/googlenet.jl

@@ -36,7 +36,7 @@ Create an Inception-v1 model (commonly referred to as GoogLeNet)
 
   - `nclasses`: the number of output classes
 """
-function googlenet(; inchannels::Integer = 3, nclasses::Integer = 1000)
+function googlenet(; dropout_rate = 0.4, inchannels::Integer = 3, nclasses::Integer = 1000)
     backbone = Chain(Conv((7, 7), inchannels => 64; stride = 2, pad = 3),
                      MaxPool((3, 3); stride = 2, pad = 1),
                      Conv((1, 1), 64 => 64),
@@ -53,7 +53,7 @@ function googlenet(; inchannels::Integer = 3, nclasses::Integer = 1000)
                      MaxPool((3, 3); stride = 2, pad = 1),
                      _inceptionblock(832, 256, 160, 320, 32, 128, 128),
                      _inceptionblock(832, 384, 192, 384, 48, 128, 128))
-    return Chain(backbone, create_classifier(1024, nclasses; dropout_rate = 0.4))
+    return Chain(backbone, create_classifier(1024, nclasses; dropout_rate))
 end
 
 """

src/convnets/inception/inceptionresnetv2.jl

@@ -96,7 +96,8 @@ function inceptionresnetv2(; inchannels::Integer = 3, dropout_rate = 0.0,
 end
 
 """
-    InceptionResNetv2(; pretrain::Bool = false, inchannels::Integer = 3, dropout_rate = 0.0, nclasses::Integer = 1000)
+    InceptionResNetv2(; pretrain::Bool = false, inchannels::Integer = 3,
+                        nclasses::Integer = 1000)
 
 Creates an InceptionResNetv2 model.
 ([reference](https://arxiv.org/abs/1602.07261))
@@ -118,9 +119,8 @@ end
 @functor InceptionResNetv2
 
 function InceptionResNetv2(; pretrain::Bool = false, inchannels::Integer = 3,
-                           dropout_rate = 0.0,
                            nclasses::Integer = 1000)
-    layers = inceptionresnetv2(; inchannels, dropout_rate, nclasses)
+    layers = inceptionresnetv2(; inchannels, nclasses)
     if pretrain
         loadpretrain!(layers, "InceptionResNetv2")
     end

src/convnets/inception/inceptionv3.jl

@@ -135,7 +135,7 @@ Create an Inception-v3 model ([reference](https://arxiv.org/abs/1512.00567v3)).
 
   - `nclasses`: the number of output classes
 """
-function inceptionv3(; inchannels::Integer = 3, nclasses::Integer = 1000)
+function inceptionv3(; dropout_rate = 0.2, inchannels::Integer = 3, nclasses::Integer = 1000)
     backbone = Chain(conv_norm((3, 3), inchannels, 32; stride = 2)...,
                      conv_norm((3, 3), 32, 32)...,
                      conv_norm((3, 3), 32, 64; pad = 1)...,
@@ -154,7 +154,7 @@ function inceptionv3(; inchannels::Integer = 3, nclasses::Integer = 1000)
                      inceptionv3_d(768),
                      inceptionv3_e(1280),
                      inceptionv3_e(2048))
-    return Chain(backbone, create_classifier(2048, nclasses; dropout_rate = 0.2))
+    return Chain(backbone, create_classifier(2048, nclasses; dropout_rate))
 end
 
 """

src/convnets/inception/xception.jl

@@ -66,8 +66,7 @@ function xception(; dropout_rate = 0.0, inchannels::Integer = 3, nclasses::Integ
                      xception_block(728, 1024, 2; stride = 2, grow_at_start = false),
                      depthwise_sep_conv_norm((3, 3), 1024, 1536; pad = 1)...,
                      depthwise_sep_conv_norm((3, 3), 1536, 2048; pad = 1)...)
-    classifier = create_classifier(2048, nclasses; dropout_rate)
-    return Chain(backbone, classifier)
+    return Chain(backbone, create_classifier(2048, nclasses; dropout_rate))
 end
 
 """

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 SHORTCUT_DICT = Dict(:A => (downsample_identity, downsample_identity),
                            :B => (downsample_conv, downsample_identity),
                            :C => (downsample_conv, downsample_conv),
                            :D => (downsample_pool, downsample_identity))
@@ -288,8 +288,8 @@ function resnet(img_dims, stem, get_layers, block_repeats::AbstractVector{<:Inte
     stage_blocks = resnet_stages(get_layers, block_repeats, connection)
     backbone = Chain(stem, stage_blocks)
     # Build the classifier head
-    nfeaturemaps = Flux.outputsize(backbone, img_dims; padbatch = true)[3]
-    classifier = classifier_fn(nfeaturemaps)
+    # nfeaturemaps = Flux.outputsize(backbone, img_dims; padbatch = true)[3]
+    classifier = classifier_fn(2048)
     return Chain(backbone, classifier)
 end
 
@@ -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, SHORTCUT_DICT[downsample_opt]; kwargs...)
 end
 
 # block-layer configurations for ResNet-like models

src/convnets/resnets/res2net.jl

@@ -22,32 +22,32 @@ Creates a bottleneck block as described in the Res2Net paper.
 """
 function bottle2neck(inplanes::Integer, planes::Integer; stride::Integer = 1,
                      cardinality::Integer = 1, base_width::Integer = 26,
-                     scale::Integer = 4, activation = relu, norm_layer = BatchNorm,
-                     revnorm::Bool = false, attn_fn = planes -> identity)
+                     scale::Integer = 4, activation = relu, is_first::Bool = false,
+                     norm_layer = BatchNorm, revnorm::Bool = false,
+                     attn_fn = planes -> identity)
     width = fld(planes * base_width, 64) * cardinality
     outplanes = planes * 4
-    is_first = stride > 1
     pool = is_first && scale > 1 ? MeanPool((3, 3); stride, pad = 1) : identity
     conv_bns = [Chain(conv_norm((3, 3), width => width, activation; norm_layer, stride,
                                 pad = 1, groups = cardinality, bias = false)...)
-                for _ in 1:(max(1, scale - 1))]
+                for _ in 1:max(1, scale - 1)]
     reslayer = is_first ? Parallel(cat_channels, pool, conv_bns...) :
-               Parallel(cat_channels, identity, PairwiseFusion(+, conv_bns...))
+               Parallel(cat_channels, identity, Chain(PairwiseFusion(+, conv_bns...)))
     tuplify(x) = is_first ? tuple(x...) : tuple(x[1], tuple(x[2:end]...))
-    return Chain(conv_norm((1, 1), inplanes => width * scale, activation;
-                           norm_layer, revnorm, bias = false)...,
-                 chunk$(; size = width, dims = 3),
-                 tuplify, reslayer,
-                 conv_norm((1, 1), width * scale => outplanes, activation;
-                           norm_layer, revnorm, bias = false)...,
-                 attn_fn(outplanes))
+    layers = [conv_norm((1, 1), inplanes => width * scale, activation;
+                        norm_layer, revnorm, bias = false)...,
+              chunk$(; size = width, dims = 3), tuplify, reslayer,
+              conv_norm((1, 1), width * scale => outplanes, activation;
+                        norm_layer, revnorm, bias = false)...,
+              attn_fn(outplanes)]
+    return Chain(filter(!=(identity), layers)...)
 end
 
 function bottle2neck_builder(block_repeats::AbstractVector{<:Integer};
                              inplanes::Integer = 64, cardinality::Integer = 1,
                              base_width::Integer = 26, scale::Integer = 4,
                              expansion::Integer = 4, norm_layer = BatchNorm,
-                             revnorm::Bool = false, activation = relu, 
+                             revnorm::Bool = false, activation = relu,
                              attn_fn = planes -> identity,
                              stride_fn = resnet_stride, planes_fn = resnet_planes,
                              downsample_tuple = (downsample_conv, downsample_identity))
@@ -63,8 +63,9 @@ function bottle2neck_builder(block_repeats::AbstractVector{<:Integer};
         stride = stride_fn(stage_idx, block_idx)
         downsample_fn = (stride != 1 || inplanes != planes * expansion) ?
                         downsample_tuple[1] : downsample_tuple[2]
+        is_first = (stride > 1 || downsample_fn != downsample_tuple[2]) ? true : false
         block = bottle2neck(inplanes, planes; stride, cardinality, base_width, scale,
-                            activation, norm_layer, revnorm, attn_fn)
+                            activation, is_first, norm_layer, revnorm, attn_fn)
         downsample = downsample_fn(inplanes, planes * expansion; stride, norm_layer,
                                    revnorm)
         return block, downsample
@@ -92,19 +93,25 @@ Creates a Res2Net model with the specified depth, scale, and base width.
 struct Res2Net
     layers::Any
 end
+@functor Res2Net
 
 function Res2Net(depth::Integer; pretrain::Bool = false, scale::Integer = 4,
                  base_width::Integer = 26, inchannels::Integer = 3,
                  nclasses::Integer = 1000)
     _checkconfig(depth, sort(collect(keys(RESNET_CONFIGS)))[3:end])
-    layers = resnet(:bottle2neck, RESNET_CONFIGS[depth][2], :C; base_width, scale,
+    layers = resnet(:bottle2neck, RESNET_CONFIGS[depth][2]; base_width, scale,
                     inchannels, nclasses)
     if pretrain
         loadpretrain!(layers, string("Res2Net", depth, "_", base_width, "x", scale))
     end
-    return ResNet(layers)
+    return Res2Net(layers)
 end
 
+(m::Res2Net)(x) = m.layers(x)
+
+backbone(m::Res2Net) = m.layers[1]
+classifier(m::Res2Net) = m.layers[2]
+
 """
     Res2NeXt(depth::Integer; pretrain::Bool = false, scale::Integer = 4,
              base_width::Integer = 4, cardinality::Integer = 8,
@@ -126,17 +133,23 @@ Creates a Res2NeXt model with the specified depth, scale, base width and cardina
 struct Res2NeXt
     layers::Any
 end
+@functor Res2NeXt
 
 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], :C; base_width, scale,
+    layers = resnet(:bottle2neck, RESNET_CONFIGS[depth][2]; base_width, scale,
                     cardinality, inchannels, nclasses)
     if pretrain
         loadpretrain!(layers,
                       string("Res2NeXt", depth, "_", base_width, "x", cardinality,
                              "x", scale))
     end
-    return ResNet(layers)
+    return Res2NeXt(layers)
 end
+
+(m::Res2NeXt)(x) = m.layers(x)
+
+backbone(m::Res2NeXt) = m.layers[1]
+classifier(m::Res2NeXt) = m.layers[2]

src/utilities.jl

@@ -38,7 +38,10 @@ Concatenate `x` and `y` (and any `z`s) along the channel dimension (third dimens
 Equivalent to `cat(x, y, zs...; dims=3)`.
 Convenient reduction operator for use with `Parallel`.
 """
-cat_channels(xy...) = cat(xy...; dims = Val(3))
+cat_channels(xs::AbstractArray...) = cat(xs...; dims = Val(3))
+cat_channels(x::AbstractArray, y::Tuple) = cat_channels(x, y...)
+cat_channels(x::Tuple, y::AbstractArray...) = cat_channels(x..., y...)
+cat_channels(x::Tuple) = cat_channels(x...)
 
 """
     swapdims(perm)

test/convnets.jl

@@ -126,7 +126,7 @@ end
     @testset for (base_width, scale) in [(26, 4), (48, 2), (14, 8), (26, 6), (26, 8)]
         m = Res2Net(50; base_width, scale)
         @test size(m(x_224)) == (1000, 1)
-        if (Res2Net, depth, cardinality, base_width) in PRETRAINED_MODELS
+        if (Res2Net, depth, base_width, scale) in PRETRAINED_MODELS
             @test acctest(Res2Net(depth, pretrain = true))
         else
             @test_throws ArgumentError Res2Net(depth, pretrain = true)
@@ -137,7 +137,7 @@ end
     @testset for (base_width, scale) in [(26, 4)]
         m = Res2Net(101; base_width, scale)
         @test size(m(x_224)) == (1000, 1)
-        if (Res2Net, depth, cardinality, base_width) in PRETRAINED_MODELS
+        if (Res2Net, depth, base_width, scale) in PRETRAINED_MODELS
             @test acctest(Res2Net(depth, pretrain = true))
         else
             @test_throws ArgumentError Res2Net(depth, pretrain = true)