Minor refactor of cnn_stages

And no, final cleanup was not final.

1. Fix `width_mult` calculations
2. Hopefully all the parameters line up now for all widths of the MobileNets
3. `MNASNet` wasn't `@functor`ed
4. Random docstring link fixes

Author: theabhirath

src/convnets/builders/core.jl

@@ -1,31 +1,19 @@
-# TODO potentially refactor other CNNs to use this
-function cnn_stages(get_layers, block_repeats::AbstractVector{<:Integer}, connection)
+function cnn_stages(get_layers, block_repeats::AbstractVector{<:Integer},
+                    connection = nothing)
     # Construct each stage
     stages = []
     for (stage_idx, nblocks) in enumerate(block_repeats)
         # Construct the blocks for each stage
         blocks = map(1:nblocks) do block_idx
             branches = get_layers(stage_idx, block_idx)
+            if isnothing(connection)
+                @assert length(branches)==1 "get_layers should return a single branch for 
+                each block if no connection is specified"
+            end
             return length(branches) == 1 ? only(branches) :
                    Parallel(connection, branches...)
         end
         push!(stages, Chain(blocks...))
     end
     return stages
 end
-
-function cnn_stages(get_layers, block_repeats::AbstractVector{<:Integer})
-    # Construct each stage
-    stages = []
-    for (stage_idx, nblocks) in enumerate(block_repeats)
-        # Construct the blocks for each stage
-        blocks = map(1:nblocks) do block_idx
-            branches = get_layers(stage_idx, block_idx)
-            @assert length(branches)==1 "get_layers should return a single branch for each
-          block if no connection is specified"
-            return only(branches)
-        end
-        push!(stages, Chain(blocks...))
-    end
-    return stages
-end

src/convnets/builders/mbconv.jl

@@ -2,7 +2,9 @@ function dwsepconv_builder(block_configs, inplanes::Integer, stage_idx::Integer,
                            width_mult::Real; norm_layer = BatchNorm, kwargs...)
     block_fn, k, outplanes, stride, nrepeats, activation = block_configs[stage_idx]
     outplanes = floor(Int, outplanes * width_mult)
-    inplanes = stage_idx == 1 ? inplanes : block_configs[stage_idx - 1][3]
+    if stage_idx != 1
+        inplanes = floor(Int, block_configs[stage_idx - 1][3] * width_mult)
+    end
     function get_layers(block_idx::Integer)
         inplanes = block_idx == 1 ? inplanes : outplanes
         stride = block_idx == 1 ? stride : 1
@@ -23,8 +25,9 @@ function mbconv_builder(block_configs, inplanes::Integer, stage_idx::Integer,
     if !isnothing(reduction)
         reduction = !se_from_explanes ? reduction * expansion : reduction
     end
-    inplanes = stage_idx == 1 ? inplanes : block_configs[stage_idx - 1][3]
-    inplanes = _round_channels(inplanes * width_mult, divisor)
+    if stage_idx != 1
+        inplanes = _round_channels(block_configs[stage_idx - 1][3] * width_mult, divisor)
+    end
     outplanes = _round_channels(outplanes * width_mult, divisor)
     function get_layers(block_idx::Integer)
         inplanes = block_idx == 1 ? inplanes : outplanes

src/convnets/efficientnets/core.jl

@@ -5,16 +5,17 @@ function efficientnet(block_configs::AbstractVector{<:Tuple}; inplanes::Integer,
                       inchannels::Integer = 3, nclasses::Integer = 1000)
     layers = []
     # stem of the model
+    inplanes = _round_channels(inplanes * scalings[1])
     append!(layers,
-            conv_norm((3, 3), inchannels, _round_channels(inplanes * scalings[1], 8),
-                      swish; norm_layer, stride = 2, pad = SamePad()))
+            conv_norm((3, 3), inchannels, inplanes, swish; norm_layer, stride = 2,
+                      pad = SamePad()))
     # building inverted residual blocks
     get_layers, block_repeats = mbconv_stack_builder(block_configs, inplanes; scalings,
                                                      norm_layer)
     append!(layers, cnn_stages(get_layers, block_repeats, +))
     # building last layers
     append!(layers,
-            conv_norm((1, 1), _round_channels(block_configs[end][3] * scalings[1], 8),
+            conv_norm((1, 1), _round_channels(block_configs[end][3] * scalings[1]),
                       headplanes, swish; pad = SamePad()))
     return Chain(Chain(layers...), create_classifier(headplanes, nclasses; dropout_rate))
 end

src/convnets/mobilenets/mnasnet.jl

@@ -25,7 +25,7 @@ function mnasnet(block_configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
     norm_layer = (args...; kwargs...) -> BatchNorm(args...; momentum = _MNASNET_BN_MOMENTUM,
                                                    kwargs...)
     # building first layer
-    inplanes = _round_channels(inplanes * width_mult, 8)
+    inplanes = _round_channels(inplanes * width_mult)
     layers = []
     append!(layers,
             conv_norm((3, 3), inchannels, inplanes, relu; stride = 2, pad = 1,
@@ -35,8 +35,8 @@ function mnasnet(block_configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
                                                      norm_layer)
     append!(layers, cnn_stages(get_layers, block_repeats, +))
     # building last layers
-    outplanes = _round_channels(block_configs[end][3] * width_mult, 8)
-    headplanes = _round_channels(max_width * max(1, width_mult), 8)
+    outplanes = _round_channels(block_configs[end][3] * width_mult)
+    headplanes = _round_channels(max_width * max(1, width_mult))
     append!(layers,
             conv_norm((1, 1), outplanes, headplanes, relu; norm_layer))
     return Chain(Chain(layers...), create_classifier(headplanes, nclasses; dropout_rate))
@@ -119,3 +119,8 @@ function MNASNet(config::Symbol; width_mult::Real = 1, pretrain::Bool = false,
     end
     return MNASNet(layers)
 end
+
+(m::MNASNet)(x) = m.layers(x)
+
+backbone(m::MNASNet) = m.layers[1]
+classifier(m::MNASNet) = m.layers[2]

src/convnets/mobilenets/mobilenetv1.jl

@@ -23,16 +23,18 @@ Create a MobileNetv1 model ([reference](https://arxiv.org/abs/1704.04861v1)).
 """
 function mobilenetv1(config::AbstractVector{<:Tuple}; width_mult::Real = 1,
                      activation = relu, dropout_rate = nothing,
-                     inchannels::Integer = 3, nclasses::Integer = 1000)
+                     inplanes::Integer = 32, inchannels::Integer = 3,
+                     nclasses::Integer = 1000)
     layers = []
     # stem of the model
+    inplanes = floor(Int, inplanes * width_mult)
     append!(layers,
-            conv_norm((3, 3), inchannels, config[1][3], activation; stride = 2, pad = 1))
+            conv_norm((3, 3), inchannels, inplanes, activation; stride = 2, pad = 1))
     # building inverted residual blocks
-    get_layers, block_repeats = mbconv_stack_builder(config, config[1][3]; width_mult)
+    get_layers, block_repeats = mbconv_stack_builder(config, inplanes; width_mult)
     append!(layers, cnn_stages(get_layers, block_repeats))
-    return Chain(Chain(layers...),
-                 create_classifier(config[end][3], nclasses; dropout_rate))
+    outplanes = floor(Int, config[end][3] * width_mult)
+    return Chain(Chain(layers...), create_classifier(outplanes, nclasses; dropout_rate))
 end
 
 # Layer configurations for MobileNetv1
@@ -45,14 +47,10 @@ end
 const MOBILENETV1_CONFIGS = [
     # f, k, c, s, n, a
     (dwsep_conv_bn, 3, 64, 1, 1, relu6),
-    (dwsep_conv_bn, 3, 128, 2, 1, relu6),
-    (dwsep_conv_bn, 3, 128, 1, 1, relu6),
-    (dwsep_conv_bn, 3, 256, 2, 1, relu6),
-    (dwsep_conv_bn, 3, 256, 1, 1, relu6),
-    (dwsep_conv_bn, 3, 512, 2, 1, relu6),
-    (dwsep_conv_bn, 3, 512, 1, 5, relu6),
-    (dwsep_conv_bn, 3, 1024, 2, 1, relu6),
-    (dwsep_conv_bn, 3, 1024, 1, 1, relu6),
+    (dwsep_conv_bn, 3, 128, 2, 2, relu6),
+    (dwsep_conv_bn, 3, 256, 2, 2, relu6),
+    (dwsep_conv_bn, 3, 512, 2, 6, relu6),
+    (dwsep_conv_bn, 3, 1024, 2, 2, relu6),
 ]
 
 """

src/convnets/mobilenets/mobilenetv2.jl

@@ -37,7 +37,7 @@ function mobilenetv2(block_configs::AbstractVector{<:Tuple}; width_mult::Real =
                                                      divisor)
     append!(layers, cnn_stages(get_layers, block_repeats, +))
     # building last layers
-    outplanes = _round_channels(block_configs[end][3], divisor)
+    outplanes = _round_channels(block_configs[end][3] * width_mult, divisor)
     headplanes = _round_channels(max_width * max(1, width_mult), divisor)
     append!(layers, conv_norm((1, 1), outplanes, headplanes, relu6))
     return Chain(Chain(layers...), create_classifier(headplanes, nclasses; dropout_rate))

src/convnets/mobilenets/mobilenetv3.jl

@@ -28,7 +28,7 @@ function mobilenetv3(configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
                      max_width::Integer = 1024, dropout_rate = 0.2,
                      inchannels::Integer = 3, nclasses::Integer = 1000)
     # building first layer
-    inplanes = _round_channels(16 * width_mult, 8)
+    inplanes = _round_channels(16 * width_mult)
     layers = []
     append!(layers,
             conv_norm((3, 3), inchannels, inplanes, hardswish; stride = 2, pad = 1))
@@ -38,12 +38,11 @@ function mobilenetv3(configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
                                                      se_round_fn = _round_channels)
     append!(layers, cnn_stages(get_layers, block_repeats, +))
     # building last layers
-    explanes = _round_channels(configs[end][3] * width_mult, 8)
-    midplanes = _round_channels(explanes * configs[end][4], 8)
-    headplanes = _round_channels(max_width * width_mult, 8)
+    explanes = _round_channels(configs[end][3] * width_mult)
+    midplanes = _round_channels(explanes * configs[end][4])
     append!(layers, conv_norm((1, 1), explanes, midplanes, hardswish))
     return Chain(Chain(layers...),
-                 create_classifier(midplanes, headplanes, nclasses,
+                 create_classifier(midplanes, max_width, nclasses,
                                    (hardswish, identity); dropout_rate))
 end
 

src/convnets/resnets/resnet.jl

@@ -2,7 +2,7 @@
     ResNet(depth::Integer; pretrain::Bool = false, inchannels::Integer = 3, nclasses::Integer = 1000)
 
 Creates a ResNet model with the specified depth.
-((reference)[https://arxiv.org/abs/1512.03385])
+([reference](https://arxiv.org/abs/1512.03385))
 
 # Arguments
 
@@ -39,7 +39,7 @@ classifier(m::ResNet) = m.layers[2]
 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])
+([reference](https://arxiv.org/abs/1605.07146))
 
 # Arguments
 

src/convnets/resnets/resnext.jl

@@ -3,7 +3,7 @@
             base_width::Integer = 4, inchannels::Integer = 3, nclasses::Integer = 1000)
 
 Creates a ResNeXt model with the specified depth, cardinality, and base width.
-((reference)[https://arxiv.org/abs/1611.05431])
+([reference](https://arxiv.org/abs/1611.05431))
 
 # Arguments
 

src/convnets/resnets/seresnet.jl

@@ -2,7 +2,7 @@
     SEResNet(depth::Integer; pretrain::Bool = false, inchannels::Integer = 3, nclasses::Integer = 1000)
 
 Creates a SEResNet model with the specified depth.
-((reference)[https://arxiv.org/pdf/1709.01507.pdf])
+([reference](https://arxiv.org/pdf/1709.01507.pdf))
 
 # Arguments
 
@@ -43,7 +43,7 @@ classifier(m::SEResNet) = m.layers[2]
               base_width::Integer = 4, inchannels::Integer = 3, nclasses::Integer = 1000)
 
 Creates a SEResNeXt model with the specified depth, cardinality, and base width.
-((reference)[https://arxiv.org/pdf/1709.01507.pdf])
+([reference](https://arxiv.org/pdf/1709.01507.pdf))
 
 # Arguments