Fix minor hiccups

Author: theabhirath

src/convnets/efficientnets/core.jl

@@ -1,7 +1,7 @@
 function mbconv_builder(block_configs::AbstractVector{NTuple{6, Int}},
                         stage_idx::Integer; scalings::NTuple{2, Real} = (1, 1),
                         norm_layer = BatchNorm)
-    depth_mult, width_mult = scalings
+    width_mult, depth_mult = scalings
     k, inplanes, outplanes, expansion, stride, nrepeats = block_configs[stage_idx]
     inplanes = _round_channels(inplanes * width_mult, 8)
     outplanes = _round_channels(outplanes * width_mult, 8)
@@ -17,7 +17,8 @@ function mbconv_builder(block_configs::AbstractVector{NTuple{6, Int}},
 end
 
 function fused_mbconv_builder(block_configs::AbstractVector{NTuple{6, Int}},
-                              stage_idx::Integer; norm_layer = BatchNorm)
+                              stage_idx::Integer; scalings::NTuple{2, Real} = (1, 1),
+                              norm_layer = BatchNorm)
     k, inplanes, outplanes, expansion, stride, nrepeats = block_configs[stage_idx]
     function get_layers(block_idx)
         inplanes = block_idx == 1 ? inplanes : outplanes
@@ -40,22 +41,22 @@ end
 
 function efficientnet(block_configs::AbstractVector{NTuple{6, Int}},
                       residual_fns::AbstractVector; scalings::NTuple{2, Real} = (1, 1),
-                      headplanes::Integer = _round_channels(block_configs[end][3] *
-                                                            scalings[2], 8) * 4,
+                      headplanes::Integer = block_configs[end][3] * 4,
                       norm_layer = BatchNorm, dropout_rate = nothing,
                       inchannels::Integer = 3, nclasses::Integer = 1000)
     layers = []
     # stem of the model
     append!(layers,
-            conv_norm((3, 3), inchannels, block_configs[1][2], swish; norm_layer,
-                      stride = 2, pad = SamePad()))
+            conv_norm((3, 3), inchannels,
+                      _round_channels(block_configs[1][2] * scalings[1], 8), swish;
+                      norm_layer, stride = 2, pad = SamePad()))
     # building inverted residual blocks
     get_layers, block_repeats = efficientnet_builder(block_configs, residual_fns;
                                                      scalings, norm_layer)
     append!(layers, resnet_stages(get_layers, block_repeats, +))
     # building last layers
     append!(layers,
-            conv_norm((1, 1), _round_channels(block_configs[end][3] * scalings[2], 8),
+            conv_norm((1, 1), _round_channels(block_configs[end][3] * scalings[1], 8),
                       headplanes, swish; pad = SamePad()))
     return Chain(Chain(layers...), create_classifier(headplanes, nclasses; dropout_rate))
 end

src/convnets/efficientnets/efficientnetv2.jl

@@ -56,7 +56,8 @@ function EfficientNetv2(config::Symbol; pretrain::Bool = false,
                         inchannels::Integer = 3, nclasses::Integer = 1000)
     _checkconfig(config, sort(collect(keys(EFFNETV2_CONFIGS))))
     layers = efficientnet(EFFNETV2_CONFIGS[config],
-                          vcat(fill(fused_mbconv_builder, 3), fill(mbconv_builder, 4));
+                          vcat(fill(fused_mbconv_builder, 3),
+                               fill(mbconv_builder, length(EFFNETV2_CONFIGS[config]) - 3));
                           headplanes = 1280, inchannels, nclasses)
     if pretrain
         loadpretrain!(layers, string("efficientnetv2"))