Refine invertedresidual

Author: theabhirath

src/convnets/efficientnet.jl

@@ -36,11 +36,12 @@ function efficientnet(scalings, block_configs;
         out_channels = _round_channels(scalew(o), 8)
         repeats = scaled(n)
         push!(blocks,
-              invertedresidual(k, in_channels, in_channels * e, out_channels, swish;
+              invertedresidual((k, k), in_channels, in_channels * e, out_channels, swish;
                                stride = s, reduction = 4))
         for _ in 1:(repeats - 1)
             push!(blocks,
-                  invertedresidual(k, out_channels, out_channels * e, out_channels, swish;
+                  invertedresidual((k, k), out_channels, out_channels * e, out_channels,
+                                   swish;
                                    stride = 1, reduction = 4))
         end
     end

src/convnets/mobilenet/mobilenetv2.jl

@@ -30,7 +30,7 @@ function mobilenetv2(width_mult, configs; inchannels = 3, max_width = 1280, ncla
         outplanes = _round_channels(c * width_mult, width_mult == 0.1 ? 4 : 8)
         for i in 1:n
             push!(layers,
-                  invertedresidual(3, inplanes, inplanes * t, outplanes, a;
+                  invertedresidual((3, 3), inplanes, inplanes * t, outplanes, a;
                                    stride = i == 1 ? s : 1))
             inplanes = outplanes
         end

src/convnets/mobilenet/mobilenetv3.jl

@@ -36,7 +36,7 @@ function mobilenetv3(width_mult, configs; inchannels = 3, max_width = 1024, ncla
         outplanes = _round_channels(c * width_mult, 8)
         explanes = _round_channels(inplanes * t, 8)
         push!(layers,
-              invertedresidual(k, inplanes, explanes, outplanes, a;
+              invertedresidual((k, k), inplanes, explanes, outplanes, a;
                                stride = s, reduction = r))
         inplanes = outplanes
     end

src/layers/conv.jl

@@ -114,23 +114,28 @@ Create a basic inverted residual block for MobileNet variants
   - `reduction`: The reduction factor for the number of hidden feature maps
     in a squeeze and excite layer (see [`squeeze_excite`](#)).
 """
-function invertedresidual(kernel_size, inplanes, hidden_planes, outplanes,
-                          activation = relu; stride, reduction = nothing)
+function invertedresidual(kernel_size, inplanes::Integer, hidden_planes::Integer,
+                          outplanes::Integer, activation = relu; stride::Integer,
+                          reduction::Union{Nothing, Integer} = nothing)
     @assert stride in [1, 2] "`stride` has to be 1 or 2"
     pad = @. (kernel_size - 1) ÷ 2
-    conv1 = (inplanes == hidden_planes) ? identity :
-            Chain(conv_norm((1, 1), inplanes, hidden_planes, activation; bias = false))
+    conv1 = (inplanes == hidden_planes) ? (identity,) :
+            conv_norm((1, 1), inplanes, hidden_planes, activation; bias = false)
     selayer = isnothing(reduction) ? identity :
               squeeze_excite(hidden_planes; reduction, activation, gate_activation = hardσ,
                              norm_layer = BatchNorm)
-    invres = Chain(conv1,
+    invres = Chain(conv1...,
                    conv_norm(kernel_size, hidden_planes, hidden_planes, activation;
                              bias = false, stride, pad = pad, groups = hidden_planes)...,
                    selayer,
                    conv_norm((1, 1), hidden_planes, outplanes, identity; bias = false)...)
     return (stride == 1 && inplanes == outplanes) ? SkipConnection(invres, +) : invres
 end
 
-function invertedresidual(kernel_size::Integer, args...; kwargs...)
-    return invertedresidual((kernel_size, kernel_size), args...; kwargs...)
+function invertedresidual(kernel_size, inplanes::Integer, outplanes::Integer,
+                          activation = relu; stride::Integer, expansion,
+                          reduction::Union{Nothing, Integer} = nothing)
+    hidden_planes = Int(inplanes * expansion)
+    return invertedresidual(kernel_size, inplanes, hidden_planes, outplanes, activation;
+                            stride, reduction)
 end