Tweaks to prepare pretrained models

Author: theabhirath

src/convnets/densenet.jl

@@ -135,7 +135,7 @@ function DenseNet(config::Integer; pretrain::Bool = false, growth_rate::Integer
     _checkconfig(config, keys(DENSENET_CONFIGS))
     model = densenet(DENSENET_CONFIGS[config]; growth_rate, reduction, inchannels, nclasses)
     if pretrain
-        loadpretrain!(model, string("DenseNet", config))
+        loadpretrain!(model, string("densenet", config))
     end
     return model
 end

src/convnets/inception/inceptionv3.jl

@@ -1,5 +1,3 @@
-## Inceptionv3
-
 """
     inceptionv3_a(inplanes, pool_proj)
 

src/convnets/resnets/core.jl

@@ -280,15 +280,13 @@ function resnet_stages(get_layers, block_repeats::AbstractVector{<:Integer}, con
 end
 
 function resnet(img_dims, stem, get_layers, block_repeats::AbstractVector{<:Integer},
-                connection,
-                classifier_fn)
+                connection, classifier_fn)
     # Build stages of the ResNet
     stage_blocks = resnet_stages(get_layers, block_repeats, connection)
     backbone = Chain(stem, stage_blocks)
-    # Build the classifier head
+    # Add classifier to the backbone
     nfeaturemaps = Flux.outputsize(backbone, img_dims; padbatch = true)[3]
-    classifier = classifier_fn(nfeaturemaps)
-    return Chain(backbone, classifier)
+    return Chain(backbone, classifier_fn(nfeaturemaps))
 end
 
 function resnet(block_type::Symbol, block_repeats::AbstractVector{<:Integer};

src/convnets/resnets/resnet.jl

@@ -27,7 +27,7 @@ function ResNet(depth::Integer; pretrain::Bool = false, inchannels::Integer = 3,
     _checkconfig(depth, keys(RESNET_CONFIGS))
     layers = resnet(RESNET_CONFIGS[depth]...; inchannels, nclasses)
     if pretrain
-        loadpretrain!(layers, string("ResNet", depth))
+        loadpretrain!(layers, string("resnet", depth))
     end
     return ResNet(layers)
 end

src/convnets/resnets/resnext.jl

@@ -32,7 +32,7 @@ function ResNeXt(depth::Integer; pretrain::Bool = false, cardinality = 32,
     _checkconfig(depth, sort(collect(keys(RESNET_CONFIGS)))[3:end])
     layers = resnet(RESNET_CONFIGS[depth]...; inchannels, nclasses, cardinality, base_width)
     if pretrain
-        loadpretrain!(layers, string("ResNeXt", depth, "_", cardinality, "x", base_width))
+        loadpretrain!(layers, string("resnext", depth, "_", cardinality, "x", base_width, "d"))
     end
     return ResNeXt(layers)
 end

src/convnets/squeezenet.jl

@@ -77,7 +77,7 @@ function SqueezeNet(; pretrain::Bool = false, inchannels::Integer = 3,
                     nclasses::Integer = 1000)
     layers = squeezenet(; inchannels, nclasses)
     if pretrain
-        loadpretrain!(layers, "SqueezeNet")
+        loadpretrain!(layers, "squeezenet")
     end
     return SqueezeNet(layers)
 end

src/layers/mlp.jl

@@ -66,16 +66,19 @@ Creates a classifier head to be used for models.
 function create_classifier(inplanes::Integer, nclasses::Integer, activation = identity;
                            use_conv::Bool = false, pool_layer = AdaptiveMeanPool((1, 1)),
                            dropout_rate = nothing)
-    # Pooling
+    # Decide whether to flatten the input or not
     flatten_in_pool = !use_conv && pool_layer !== identity
     if use_conv
         @assert pool_layer === identity
         "`pool_layer` must be identity if `use_conv` is true"
     end
-    global_pool = flatten_in_pool ? [pool_layer, MLUtils.flatten] : [pool_layer]
+    classifier = []
+    flatten_in_pool ? push!(classifier, pool_layer, MLUtils.flatten) : 
+        push!(classifier, pool_layer)
+    # Dropout is applied after the pooling layer
+    isnothing(dropout_rate) ? nothing : push!(classifier, Dropout(dropout_rate))
     # Fully-connected layer
-    fc = use_conv ? Conv((1, 1), inplanes => nclasses, activation) :
-         Dense(inplanes => nclasses, activation)
-    drop = isnothing(dropout_rate) ? [] : [Dropout(dropout_rate)]
-    return Chain(global_pool..., drop..., fc)
+    use_conv ? push!(classifier, Conv((1, 1), inplanes => nclasses, activation)) :
+        push!(classifier, Dense(inplanes => nclasses, activation))
+    return Chain(classifier...)
 end

src/utilities.jl

@@ -73,6 +73,5 @@ end
 
 # Utility function for depth and configuration checks in models
 function _checkconfig(config, configs)
-    @assert config in configs
-    return "Invalid configuration. Must be one of $(sort(collect(configs)))."
+    @assert config in configs "Invalid configuration. Must be one of $(sort(collect(configs)))."
 end

test/convnets.jl

@@ -22,15 +22,14 @@ end
 
 @testset "ResNet" begin
     # Tests for pretrained ResNets
-    ## TODO: find a way to port pretrained models to the new ResNet API
     @testset "ResNet($sz)" for sz in [18, 34, 50, 101, 152]
         m = ResNet(sz)
         @test size(m(x_224)) == (1000, 1)
-        # if (ResNet, sz) in PRETRAINED_MODELS
-        #     @test acctest(ResNet(sz, pretrain = true))
-        # else
-        #     @test_throws ArgumentError ResNet(sz, pretrain = true)
-        # end
+        if (ResNet, sz) in PRETRAINED_MODELS
+            @test acctest(ResNet(sz, pretrain = true))
+        else
+            @test_throws ArgumentError ResNet(sz, pretrain = true)
+        end
     end
 
     @testset "resnet" begin
@@ -79,9 +78,9 @@ end
                 m = ResNeXt(depth; cardinality, base_width)
                 @test size(m(x_224)) == (1000, 1)
                 if (ResNeXt, depth, cardinality, base_width) in PRETRAINED_MODELS
-                    @test acctest(ResNeXt(depth, pretrain = true))
+                    @test acctest(ResNeXt(depth; cardinality, base_width, pretrain = true))
                 else
-                    @test_throws ArgumentError ResNeXt(depth, pretrain = true)
+                    @test_throws ArgumentError ResNeXt(depth; cardinality, base_width, pretrain = true)
                 end
                 @test gradtest(m, x_224)
                 _gc()

test/runtests.jl

@@ -8,11 +8,21 @@ const PRETRAINED_MODELS = [
     (VGG, 13, false),
     (VGG, 16, false),
     (VGG, 19, false),
+    SqueezeNet,
     (ResNet, 18),
     (ResNet, 34),
     (ResNet, 50),
     (ResNet, 101),
     (ResNet, 152),
+    (WideResNet, 50),
+    (WideResNet, 101),
+    (ResNeXt, 50, 32, 4),
+    (ResNeXt, 101, 64, 4),
+    (ResNeXt, 101, 32, 8),
+    (DenseNet, 121),
+    (DenseNet, 161),
+    (DenseNet, 169),
+    (DenseNet, 201),
 ]
 
 function _gc()