Refactor: mbconv instead of invertedresidual

Also fix bug in EfficientNet models

Author: theabhirath

src/convnets/efficientnet/efficientnet.jl

@@ -17,35 +17,37 @@ Create an EfficientNet model ([reference](https://arxiv.org/abs/1905.11946v5)).
       + `e`: expansion ratio
       + `i`: block input channels (will be scaled by global width scaling)
       + `o`: block output channels (will be scaled by global width scaling)
-  - `max_width`: The maximum number of feature maps in any layer of the network
   - `inchannels`: number of input channels
   - `nclasses`: number of output classes
 """
 function efficientnet(scalings::NTuple{2, Real},
                       block_configs::AbstractVector{NTuple{6, Int}};
-                      max_width::Integer = 1280, inchannels::Integer = 3,
-                      nclasses::Integer = 1000)
+                      inchannels::Integer = 3, nclasses::Integer = 1000)
+    # building first layer
     wscale, dscale = scalings
     scalew(w) = wscale ≈ 1 ? w : ceil(Int64, wscale * w)
     scaled(d) = dscale ≈ 1 ? d : ceil(Int64, dscale * d)
     outplanes = _round_channels(scalew(32), 8)
     stem = conv_norm((3, 3), inchannels, outplanes, swish; bias = false, stride = 2,
                      pad = SamePad())
+    # building inverted residual blocks
     blocks = []
     for (n, k, s, e, i, o) in block_configs
         inchannels = _round_channels(scalew(i), 8)
+        explanes = _round_channels(inchannels * e, 8)
         outplanes = _round_channels(scalew(o), 8)
         repeats = scaled(n)
         push!(blocks,
-              invertedresidual((k, k), inchannels, outplanes, swish; expansion = e,
-                               stride = s, reduction = 4))
+              mbconv((k, k), inchannels, explanes, outplanes, swish;
+                     stride = s, reduction = 4))
         for _ in 1:(repeats - 1)
             push!(blocks,
-                  invertedresidual((k, k), outplanes, outplanes, swish; expansion = e,
-                                   stride = 1, reduction = 4))
+                  mbconv((k, k), outplanes, explanes, outplanes, swish;
+                         stride = 1, reduction = 4))
         end
     end
-    headplanes = _round_channels(max_width, 8)
+    # building last layers
+    headplanes = outplanes * 4
     append!(blocks,
             conv_norm((1, 1), outplanes, headplanes, swish; bias = false, pad = SamePad()))
     return Chain(Chain(stem..., blocks...), create_classifier(headplanes, nclasses))

src/convnets/efficientnet/efficientnetv2.jl

@@ -32,12 +32,19 @@ function efficientnetv2(config::AbstractVector{<:Tuple}; max_width::Integer = 17
             conv_norm((3, 3), inchannels, inplanes, swish; pad = 1, stride = 2,
                       bias = false))
     # building inverted residual blocks
-    for (t, c, n, s, reduction) in config
-        outplanes = _round_channels(c * width_mult, 8)
+    for (t, inplanes, outplanes, n, s, reduction) in config
+        explanes = _round_channels(inplanes * t, 8)
         for i in 1:n
-            push!(layers,
-                  invertedresidual((3, 3), inplanes, outplanes, swish; expansion = t,
-                                   stride = i == 1 ? s : 1, reduction))
+            stride = i == 1 ? s : 1
+            if isnothing(reduction)
+                push!(layers,
+                      fused_mbconv((3, 3), inplanes, explanes, outplanes, swish; stride))
+            else
+                inplanes = _round_channels(inplanes * width_mult, 8)
+                outplanes = _round_channels(outplanes * width_mult, 8)
+                push!(layers,
+                      mbconv((3, 3), inplanes, explanes, outplanes, swish; stride))
+            end
             inplanes = outplanes
         end
     end
@@ -49,37 +56,37 @@ function efficientnetv2(config::AbstractVector{<:Tuple}; max_width::Integer = 17
 end
 
 # block configs for EfficientNetv2
-const EFFNETV2_CONFIGS = Dict(:small => [# t, c, n, s, r
-                                  (1, 24, 2, 1, nothing),
-                                  (4, 48, 4, 2, nothing),
-                                  (4, 64, 4, 2, nothing),
-                                  (4, 128, 6, 2, 4),
-                                  (6, 160, 9, 1, 4),
-                                  (6, 256, 15, 2, 4)],
-                              :medium => [# t, c, n, s, r
-                                  (1, 24, 3, 1, nothing),
-                                  (4, 48, 5, 2, nothing),
-                                  (4, 80, 5, 2, nothing),
-                                  (4, 160, 7, 2, 4),
-                                  (6, 176, 14, 1, 4),
-                                  (6, 304, 18, 2, 4),
-                                  (6, 512, 5, 1, 4)],
-                              :large => [# t, c, n, s, r
-                                  (1, 32, 4, 1, nothing),
-                                  (4, 64, 8, 2, nothing),
-                                  (4, 96, 8, 2, nothing),
-                                  (4, 192, 16, 2, 4),
-                                  (6, 256, 24, 1, 4),
-                                  (6, 512, 32, 2, 4),
-                                  (6, 640, 8, 1, 4)],
-                              :xlarge => [# t, c, n, s, r
-                                  (1, 32, 4, 1, nothing),
-                                  (4, 64, 8, 2, nothing),
-                                  (4, 96, 8, 2, nothing),
-                                  (4, 192, 16, 2, 4),
-                                  (6, 256, 24, 1, 4),
-                                  (6, 512, 32, 2, 4),
-                                  (6, 640, 8, 1, 4)])
+const EFFNETV2_CONFIGS = Dict(:small => [
+                                  (1, 24, 24, 2, 1, nothing),
+                                  (4, 24, 48, 4, 2, nothing),
+                                  (4, 48, 64, 4, 2, nothing),
+                                  (4, 64, 128, 6, 2, 4),
+                                  (6, 128, 160, 9, 1, 4),
+                                  (6, 160, 256, 15, 2, 4)],
+                              :medium => [
+                                  (1, 24, 24, 3, 1, nothing),
+                                  (4, 24, 48, 5, 2, nothing),
+                                  (4, 48, 80, 5, 2, nothing),
+                                  (4, 80, 160, 7, 2, 4),
+                                  (6, 160, 176, 14, 1, 4),
+                                  (6, 176, 304, 18, 2, 4),
+                                  (6, 304, 512, 5, 1, 4)],
+                              :large => [
+                                  (1, 32, 32, 4, 1, nothing),
+                                  (4, 32, 64, 7, 2, nothing),
+                                  (4, 64, 96, 7, 2, nothing),
+                                  (4, 96, 192, 10, 2, 4),
+                                  (6, 192, 224, 19, 1, 4),
+                                  (6, 224, 384, 25, 2, 4),
+                                  (6, 384, 640, 7, 1, 4)],
+                              :xlarge => [
+                                  (1, 32, 32, 4, 1, nothing),
+                                  (4, 32, 64, 8, 2, nothing),
+                                  (4, 64, 96, 8, 2, nothing),
+                                  (4, 96, 192, 16, 2, 4),
+                                  (6, 192, 256, 24, 1, 4),
+                                  (6, 256, 512, 32, 2, 4),
+                                  (6, 512, 640, 8, 1, 4)])
 
 """
     EfficientNetv2(config::Symbol; pretrain::Bool = false, width_mult::Real = 1,

src/convnets/inception/inceptionresnetv2.jl

@@ -1,64 +1,64 @@
 function mixed_5b()
-    branch1 = Chain(conv_norm((1, 1), 192, 96)...)
-    branch2 = Chain(conv_norm((1, 1), 192, 48)...,
-                    conv_norm((5, 5), 48, 64; pad = 2)...)
-    branch3 = Chain(conv_norm((1, 1), 192, 64)...,
-                    conv_norm((3, 3), 64, 96; pad = 1)...,
-                    conv_norm((3, 3), 96, 96; pad = 1)...)
+    branch1 = Chain(basic_conv_bn((1, 1), 192, 96)...)
+    branch2 = Chain(basic_conv_bn((1, 1), 192, 48)...,
+                    basic_conv_bn((5, 5), 48, 64; pad = 2)...)
+    branch3 = Chain(basic_conv_bn((1, 1), 192, 64)...,
+                    basic_conv_bn((3, 3), 64, 96; pad = 1)...,
+                    basic_conv_bn((3, 3), 96, 96; pad = 1)...)
     branch4 = Chain(MeanPool((3, 3); pad = 1, stride = 1),
-                    conv_norm((1, 1), 192, 64)...)
+                    basic_conv_bn((1, 1), 192, 64)...)
     return Parallel(cat_channels, branch1, branch2, branch3, branch4)
 end
 
 function block35(scale = 1.0f0)
-    branch1 = Chain(conv_norm((1, 1), 320, 32)...)
-    branch2 = Chain(conv_norm((1, 1), 320, 32)...,
-                    conv_norm((3, 3), 32, 32; pad = 1)...)
-    branch3 = Chain(conv_norm((1, 1), 320, 32)...,
-                    conv_norm((3, 3), 32, 48; pad = 1)...,
-                    conv_norm((3, 3), 48, 64; pad = 1)...)
-    branch4 = Chain(conv_norm((1, 1), 128, 320)...)
+    branch1 = Chain(basic_conv_bn((1, 1), 320, 32)...)
+    branch2 = Chain(basic_conv_bn((1, 1), 320, 32)...,
+                    basic_conv_bn((3, 3), 32, 32; pad = 1)...)
+    branch3 = Chain(basic_conv_bn((1, 1), 320, 32)...,
+                    basic_conv_bn((3, 3), 32, 48; pad = 1)...,
+                    basic_conv_bn((3, 3), 48, 64; pad = 1)...)
+    branch4 = Chain(basic_conv_bn((1, 1), 128, 320)...)
     return SkipConnection(Chain(Parallel(cat_channels, branch1, branch2, branch3),
                                 branch4, inputscale(scale; activation = relu)), +)
 end
 
 function mixed_6a()
-    branch1 = Chain(conv_norm((3, 3), 320, 384; stride = 2)...)
-    branch2 = Chain(conv_norm((1, 1), 320, 256)...,
-                    conv_norm((3, 3), 256, 256; pad = 1)...,
-                    conv_norm((3, 3), 256, 384; stride = 2)...)
+    branch1 = Chain(basic_conv_bn((3, 3), 320, 384; stride = 2)...)
+    branch2 = Chain(basic_conv_bn((1, 1), 320, 256)...,
+                    basic_conv_bn((3, 3), 256, 256; pad = 1)...,
+                    basic_conv_bn((3, 3), 256, 384; stride = 2)...)
     branch3 = MaxPool((3, 3); stride = 2)
     return Parallel(cat_channels, branch1, branch2, branch3)
 end
 
 function block17(scale = 1.0f0)
-    branch1 = Chain(conv_norm((1, 1), 1088, 192)...)
-    branch2 = Chain(conv_norm((1, 1), 1088, 128)...,
-                    conv_norm((7, 1), 128, 160; pad = (3, 0))...,
-                    conv_norm((1, 7), 160, 192; pad = (0, 3))...)
-    branch3 = Chain(conv_norm((1, 1), 384, 1088)...)
+    branch1 = Chain(basic_conv_bn((1, 1), 1088, 192)...)
+    branch2 = Chain(basic_conv_bn((1, 1), 1088, 128)...,
+                    basic_conv_bn((7, 1), 128, 160; pad = (3, 0))...,
+                    basic_conv_bn((1, 7), 160, 192; pad = (0, 3))...)
+    branch3 = Chain(basic_conv_bn((1, 1), 384, 1088)...)
     return SkipConnection(Chain(Parallel(cat_channels, branch1, branch2),
                                 branch3, inputscale(scale; activation = relu)), +)
 end
 
 function mixed_7a()
-    branch1 = Chain(conv_norm((1, 1), 1088, 256)...,
-                    conv_norm((3, 3), 256, 384; stride = 2)...)
-    branch2 = Chain(conv_norm((1, 1), 1088, 256)...,
-                    conv_norm((3, 3), 256, 288; stride = 2)...)
-    branch3 = Chain(conv_norm((1, 1), 1088, 256)...,
-                    conv_norm((3, 3), 256, 288; pad = 1)...,
-                    conv_norm((3, 3), 288, 320; stride = 2)...)
+    branch1 = Chain(basic_conv_bn((1, 1), 1088, 256)...,
+                    basic_conv_bn((3, 3), 256, 384; stride = 2)...)
+    branch2 = Chain(basic_conv_bn((1, 1), 1088, 256)...,
+                    basic_conv_bn((3, 3), 256, 288; stride = 2)...)
+    branch3 = Chain(basic_conv_bn((1, 1), 1088, 256)...,
+                    basic_conv_bn((3, 3), 256, 288; pad = 1)...,
+                    basic_conv_bn((3, 3), 288, 320; stride = 2)...)
     branch4 = MaxPool((3, 3); stride = 2)
     return Parallel(cat_channels, branch1, branch2, branch3, branch4)
 end
 
 function block8(scale = 1.0f0; activation = identity)
-    branch1 = Chain(conv_norm((1, 1), 2080, 192)...)
-    branch2 = Chain(conv_norm((1, 1), 2080, 192)...,
-                    conv_norm((3, 1), 192, 224; pad = (1, 0))...,
-                    conv_norm((1, 3), 224, 256; pad = (0, 1))...)
-    branch3 = Chain(conv_norm((1, 1), 448, 2080)...)
+    branch1 = Chain(basic_conv_bn((1, 1), 2080, 192)...)
+    branch2 = Chain(basic_conv_bn((1, 1), 2080, 192)...,
+                    basic_conv_bn((3, 1), 192, 224; pad = (1, 0))...,
+                    basic_conv_bn((1, 3), 224, 256; pad = (0, 1))...)
+    branch3 = Chain(basic_conv_bn((1, 1), 448, 2080)...)
     return SkipConnection(Chain(Parallel(cat_channels, branch1, branch2),
                                 branch3, inputscale(scale; activation)), +)
 end
@@ -77,12 +77,12 @@ Creates an InceptionResNetv2 model.
 """
 function inceptionresnetv2(; dropout_rate = 0.0, 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)...,
+    backbone = Chain(basic_conv_bn((3, 3), inchannels, 32; stride = 2)...,
+                     basic_conv_bn((3, 3), 32, 32)...,
+                     basic_conv_bn((3, 3), 32, 64; pad = 1)...,
                      MaxPool((3, 3); stride = 2),
-                     conv_norm((3, 3), 64, 80)...,
-                     conv_norm((3, 3), 80, 192)...,
+                     basic_conv_bn((3, 3), 64, 80)...,
+                     basic_conv_bn((3, 3), 80, 192)...,
                      MaxPool((3, 3); stride = 2),
                      mixed_5b(),
                      [block35(0.17f0) for _ in 1:10]...,
@@ -91,7 +91,7 @@ function inceptionresnetv2(; dropout_rate = 0.0, inchannels::Integer = 3,
                      mixed_7a(),
                      [block8(0.20f0) for _ in 1:9]...,
                      block8(; activation = relu),
-                     conv_norm((1, 1), 2080, 1536)...)
+                     basic_conv_bn((1, 1), 2080, 1536)...)
     return Chain(backbone, create_classifier(1536, nclasses; dropout_rate))
 end
 

src/convnets/mobilenet/mobilenetv2.jl

@@ -37,8 +37,8 @@ function mobilenetv2(configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
         outplanes = _round_channels(c * width_mult, divisor)
         for i in 1:n
             push!(layers,
-                  invertedresidual((3, 3), inplanes, outplanes, a; expansion = t,
-                                   stride = i == 1 ? s : 1))
+                  mbconv((3, 3), inplanes, round(Int, inplanes * t), outplanes, a;
+                         stride = i == 1 ? s : 1))
             inplanes = outplanes
         end
     end

src/convnets/mobilenet/mobilenetv3.jl

@@ -24,8 +24,8 @@ Create a MobileNetv3 model.
   - `nclasses`: the number of output classes
 """
 function mobilenetv3(configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
-                     max_width::Integer = 1024, inchannels::Integer = 3,
-                     nclasses::Integer = 1000)
+                     max_width::Integer = 1024, dropout_rate = 0.2,
+                     inchannels::Integer = 3, nclasses::Integer = 1000)
     # building first layer
     inplanes = _round_channels(16 * width_mult, 8)
     layers = []
@@ -39,8 +39,8 @@ function mobilenetv3(configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
         outplanes = _round_channels(c * width_mult, 8)
         explanes = _round_channels(inplanes * t, 8)
         push!(layers,
-              invertedresidual((k, k), inplanes, explanes, outplanes, activation;
-                               stride, reduction))
+              mbconv((k, k), inplanes, explanes, outplanes, activation;
+                     stride, reduction))
         inplanes = outplanes
     end
     # building last layers
@@ -49,7 +49,7 @@ function mobilenetv3(configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
     append!(layers, conv_norm((1, 1), inplanes, explanes, hardswish; bias = false))
     classifier = Chain(AdaptiveMeanPool((1, 1)), MLUtils.flatten,
                        Dense(explanes, headplanes, hardswish),
-                       Dropout(0.2),
+                       Dropout(dropout_rate),
                        Dense(headplanes, nclasses))
     return Chain(Chain(layers...), classifier)
 end

src/layers/Layers.jl

@@ -19,7 +19,7 @@ include("attention.jl")
 export MHAttention
 
 include("conv.jl")
-export conv_norm, basic_conv_bn, dwsep_conv_bn, invertedresidual
+export conv_norm, basic_conv_bn, dwsep_conv_bn, mbconv, fused_mbconv
 
 include("drop.jl")
 export DropBlock, DropPath