Refactor EfficientNets

Author: theabhirath

src/Metalhead.jl

@@ -28,18 +28,19 @@ include("convnets/resnets/resnext.jl")
 include("convnets/resnets/seresnet.jl")
 include("convnets/resnets/res2net.jl")
 ## Inceptions
-include("convnets/inception/googlenet.jl")
-include("convnets/inception/inceptionv3.jl")
-include("convnets/inception/inceptionv4.jl")
-include("convnets/inception/inceptionresnetv2.jl")
-include("convnets/inception/xception.jl")
+include("convnets/inceptions/googlenet.jl")
+include("convnets/inceptions/inceptionv3.jl")
+include("convnets/inceptions/inceptionv4.jl")
+include("convnets/inceptions/inceptionresnetv2.jl")
+include("convnets/inceptions/xception.jl")
 ## EfficientNets
-include("convnets/efficientnet/efficientnet.jl")
-include("convnets/efficientnet/efficientnetv2.jl")
+include("convnets/efficientnets/core.jl")
+include("convnets/efficientnets/efficientnet.jl")
+include("convnets/efficientnets/efficientnetv2.jl")
 ## MobileNets
-include("convnets/mobilenet/mobilenetv1.jl")
-include("convnets/mobilenet/mobilenetv2.jl")
-include("convnets/mobilenet/mobilenetv3.jl")
+include("convnets/mobilenets/mobilenetv1.jl")
+include("convnets/mobilenets/mobilenetv2.jl")
+include("convnets/mobilenets/mobilenetv3.jl")
 ## Others
 include("convnets/densenet.jl")
 include("convnets/squeezenet.jl")

src/convnets/convmixer.jl

@@ -1,5 +1,5 @@
 """
-    convmixer(planes::Integer, depth::Integer; kernel_size = (9, 9),
+    convmixer(planes::Integer, depth::Integer; kernel_size::Dims{2} = (9, 9),
               patch_size::Dims{2} = (7, 7), activation = gelu,
               inchannels::Integer = 3, nclasses::Integer = 1000)
 
@@ -16,7 +16,7 @@ Creates a ConvMixer model.
   - `inchannels`: number of input channels
   - `nclasses`: number of classes in the output
 """
-function convmixer(planes::Integer, depth::Integer; kernel_size = (9, 9),
+function convmixer(planes::Integer, depth::Integer; kernel_size::Dims{2} = (9, 9),
                    patch_size::Dims{2} = (7, 7), activation = gelu,
                    inchannels::Integer = 3, nclasses::Integer = 1000)
     stem = conv_norm(patch_size, inchannels, planes, activation; preact = true,

src/convnets/densenet.jl

@@ -12,10 +12,10 @@ Create a Densenet bottleneck layer
 """
 function dense_bottleneck(inplanes::Integer, outplanes::Integer; expansion::Integer = 4)
     inner_channels = expansion * outplanes
-    return SkipConnection(Chain(conv_norm((1, 1), inplanes, inner_channels; bias = false,
+    return SkipConnection(Chain(conv_norm((1, 1), inplanes, inner_channels;
                                           revnorm = true)...,
                                 conv_norm((3, 3), inner_channels, outplanes; pad = 1,
-                                          bias = false, revnorm = true)...),
+                                          revnorm = true)...),
                           cat_channels)
 end
 
@@ -31,7 +31,7 @@ Create a DenseNet transition sequence
   - `outplanes`: number of output feature maps
 """
 function transition(inplanes::Integer, outplanes::Integer)
-    return Chain(conv_norm((1, 1), inplanes, outplanes; bias = false, revnorm = true)...,
+    return Chain(conv_norm((1, 1), inplanes, outplanes; revnorm = true)...,
                  MeanPool((2, 2)))
 end
 
@@ -72,7 +72,7 @@ function densenet(inplanes::Integer, growth_rates; reduction = 0.5, inchannels::
                   nclasses::Integer = 1000)
     layers = []
     append!(layers,
-            conv_norm((7, 7), inchannels, inplanes; stride = 2, pad = (3, 3), bias = false))
+            conv_norm((7, 7), inchannels, inplanes; stride = 2, pad = (3, 3)))
     push!(layers, MaxPool((3, 3); stride = 2, pad = (1, 1)))
     outplanes = 0
     for (i, rates) in enumerate(growth_rates)

src/convnets/efficientnet/efficientnet.jl

@@ -0,0 +1,78 @@
+abstract type _MBConfig end
+
+struct MBConvConfig <: _MBConfig
+    kernel_size::Dims{2}
+    inplanes::Integer
+    outplanes::Integer
+    expansion::Number
+    stride::Integer
+    nrepeats::Integer
+end
+function MBConvConfig(kernel_size::Integer, inplanes::Integer, outplanes::Integer,
+                      expansion::Number, stride::Integer, nrepeats::Integer,
+                      width_mult::Number = 1, depth_mult::Number = 1)
+    inplanes = _round_channels(inplanes * width_mult, 8)
+    outplanes = _round_channels(outplanes * width_mult, 8)
+    nrepeats = ceil(Int, nrepeats * depth_mult)
+    return MBConvConfig((kernel_size, kernel_size), inplanes, outplanes, expansion,
+                        stride, nrepeats)
+end
+
+function efficientnetblock(m::MBConvConfig, norm_layer)
+    layers = []
+    explanes = _round_channels(m.inplanes * m.expansion, 8)
+    push!(layers,
+          mbconv(m.kernel_size, m.inplanes, explanes, m.outplanes, swish; norm_layer,
+                 stride = m.stride, reduction = 4))
+    explanes = _round_channels(m.outplanes * m.expansion, 8)
+    append!(layers,
+            [mbconv(m.kernel_size, m.outplanes, explanes, m.outplanes, swish; norm_layer,
+                    stride = 1, reduction = 4) for _ in 1:(m.nrepeats - 1)])
+    return Chain(layers...)
+end
+
+struct FusedMBConvConfig <: _MBConfig
+    kernel_size::Dims{2}
+    inplanes::Integer
+    outplanes::Integer
+    expansion::Number
+    stride::Integer
+    nrepeats::Integer
+end
+function FusedMBConvConfig(kernel_size::Integer, inplanes::Integer, outplanes::Integer,
+                           expansion::Number, stride::Integer, nrepeats::Integer)
+    return FusedMBConvConfig((kernel_size, kernel_size), inplanes, outplanes, expansion,
+                             stride, nrepeats)
+end
+
+function efficientnetblock(m::FusedMBConvConfig, norm_layer)
+    layers = []
+    explanes = _round_channels(m.inplanes * m.expansion, 8)
+    push!(layers,
+          fused_mbconv(m.kernel_size, m.inplanes, explanes, m.outplanes, swish;
+                       norm_layer, stride = m.stride))
+    explanes = _round_channels(m.outplanes * m.expansion, 8)
+    append!(layers,
+            [fused_mbconv(m.kernel_size, m.outplanes, explanes, m.outplanes, swish;
+                          norm_layer, stride = 1) for _ in 1:(m.nrepeats - 1)])
+    return Chain(layers...)
+end
+
+function efficientnet(block_configs::AbstractVector{<:_MBConfig};
+                      headplanes::Union{Nothing, Integer} = nothing,
+                      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].inplanes, swish; norm_layer,
+                      stride = 2, pad = SamePad()))
+    # building inverted residual blocks
+    append!(layers, [efficientnetblock(cfg, norm_layer) for cfg in block_configs])
+    # building last layers
+    outplanes = block_configs[end].outplanes
+    headplanes = isnothing(headplanes) ? outplanes * 4 : headplanes
+    append!(layers,
+            conv_norm((1, 1), outplanes, headplanes, swish; pad = SamePad()))
+    return Chain(Chain(layers...), create_classifier(headplanes, nclasses; dropout_rate))
+end