Break classifier and backbone into two sub-Chains

src/alexnet.jl

@@ -8,21 +8,21 @@ Create an AlexNet model
 - `nclasses`: the number of output classes
 """
 function alexnet(; nclasses = 1000)
-  layers = Chain(Conv((11, 11), 3 => 64, stride = (4, 4), relu, pad = (2, 2)),
-                 MaxPool((3, 3), stride = (2, 2)),
-                 Conv((5, 5), 64 => 192, relu, pad = (2, 2)),
-                 MaxPool((3, 3), stride = (2, 2)),
-                 Conv((3, 3), 192 => 384, relu, pad = (1, 1)),
-                 Conv((3, 3), 384 => 256, relu, pad = (1, 1)),
-                 Conv((3, 3), 256 => 256, relu, pad = (1, 1)),
-                 MaxPool((3, 3), stride = (2, 2)),
-                 AdaptiveMeanPool((6,6)),
-                 flatten,
-                 Dropout(0.5),
-                 Dense(256 * 6 * 6, 4096, relu),
-                 Dropout(0.5),
-                 Dense(4096, 4096, relu),
-                 Dense(4096, nclasses))
+  layers = Chain(Chain(Conv((11, 11), 3 => 64, stride = (4, 4), relu, pad = (2, 2)),
+                       MaxPool((3, 3), stride = (2, 2)),
+                       Conv((5, 5), 64 => 192, relu, pad = (2, 2)),
+                       MaxPool((3, 3), stride = (2, 2)),
+                       Conv((3, 3), 192 => 384, relu, pad = (1, 1)),
+                       Conv((3, 3), 384 => 256, relu, pad = (1, 1)),
+                       Conv((3, 3), 256 => 256, relu, pad = (1, 1)),
+                       MaxPool((3, 3), stride = (2, 2)),
+                       AdaptiveMeanPool((6,6))),
+                 Chain(flatten,
+                       Dropout(0.5),
+                       Dense(256 * 6 * 6, 4096, relu),
+                       Dropout(0.5),
+                       Dense(4096, 4096, relu),
+                       Dense(4096, nclasses)))
 
   return layers
 end

src/densenet.jl

@@ -73,10 +73,10 @@ function densenet(inplanes, growth_rates; reduction = 0.5, nclasses = 1000)
   end
   push!(layers, BatchNorm(outplanes, relu))
 
-  return Chain(layers...,
-               AdaptiveMeanPool((1, 1)),
-               flatten,
-               Dense(outplanes, nclasses))
+  return Chain(Chain(layers...),
+               Chain(AdaptiveMeanPool((1, 1)),
+                     flatten,
+                     Dense(outplanes, nclasses)))
 end
 
 """

src/googlenet.jl

@@ -39,26 +39,26 @@ Create an Inception-v1 model (commonly referred to as GoogLeNet)
 - `nclasses`: the number of output classes
 """
 function googlenet(; nclasses = 1000)
-  layers = Chain(Conv((7, 7), 3 => 64; stride = 2, pad = 3),
-                 MaxPool((3, 3), stride = 2, pad = 1),
-                 Conv((1, 1), 64 => 64),
-                 Conv((3, 3), 64 => 192; pad = 1),
-                 MaxPool((3, 3), stride = 2, pad = 1),
-                 _inceptionblock(192, 64, 96, 128, 16, 32, 32),
-                 _inceptionblock(256, 128, 128, 192, 32, 96, 64),
-                 MaxPool((3, 3), stride = 2, pad = 1),
-                 _inceptionblock(480, 192, 96, 208, 16, 48, 64),
-                 _inceptionblock(512, 160, 112, 224, 24, 64, 64),
-                 _inceptionblock(512, 128, 128, 256, 24, 64, 64),
-                 _inceptionblock(512, 112, 144, 288, 32, 64, 64),
-                 _inceptionblock(528, 256, 160, 320, 32, 128, 128),
-                 MaxPool((3, 3), stride = 2, pad = 1),
-                 _inceptionblock(832, 256, 160, 320, 32, 128, 128),
-                 _inceptionblock(832, 384, 192, 384, 48, 128, 128),
-                 AdaptiveMeanPool((1, 1)),
-                 flatten,
-                 Dropout(0.4),
-                 Dense(1024, nclasses))
+  layers = Chain(Chain(Conv((7, 7), 3 => 64; stride = 2, pad = 3),
+                       MaxPool((3, 3), stride = 2, pad = 1),
+                       Conv((1, 1), 64 => 64),
+                       Conv((3, 3), 64 => 192; pad = 1),
+                       MaxPool((3, 3), stride = 2, pad = 1),
+                       _inceptionblock(192, 64, 96, 128, 16, 32, 32),
+                       _inceptionblock(256, 128, 128, 192, 32, 96, 64),
+                       MaxPool((3, 3), stride = 2, pad = 1),
+                       _inceptionblock(480, 192, 96, 208, 16, 48, 64),
+                       _inceptionblock(512, 160, 112, 224, 24, 64, 64),
+                       _inceptionblock(512, 128, 128, 256, 24, 64, 64),
+                       _inceptionblock(512, 112, 144, 288, 32, 64, 64),
+                       _inceptionblock(528, 256, 160, 320, 32, 128, 128),
+                       MaxPool((3, 3), stride = 2, pad = 1),
+                       _inceptionblock(832, 256, 160, 320, 32, 128, 128),
+                       _inceptionblock(832, 384, 192, 384, 48, 128, 128),
+                       AdaptiveMeanPool((1, 1))),
+                 Chain(flatten,
+                       Dropout(0.4),
+                       Dense(1024, nclasses)))
 
   return layers
 end

src/inception.jl

@@ -150,28 +150,28 @@ Create an Inception-v3 model ([reference](https://arxiv.org/abs/1512.00567v3)).
     `inception3` does not currently support pretrained weights.
 """
 function inception3(; nclasses = 1000)
-  layer = Chain(conv_bn((3, 3), 3, 32; stride = 2)...,
-                conv_bn((3, 3), 32, 32)...,
-                conv_bn((3, 3), 32, 64; pad = 1)...,
-                MaxPool((3, 3), stride = 2),
-                conv_bn((1, 1), 64, 80)...,
-                conv_bn((3, 3), 80, 192)...,
-                MaxPool((3, 3), stride = 2),
-                inception_a(192, 32),
-                inception_a(256, 64),
-                inception_a(288, 64),
-                inception_b(288),
-                inception_c(768, 128),
-                inception_c(768, 160),
-                inception_c(768, 160),
-                inception_c(768, 192),
-                inception_d(768),
-                inception_e(1280),
-                inception_e(2048),
-                AdaptiveMeanPool((1, 1)),
-                Dropout(0.2),
-                flatten,
-                Dense(2048, nclasses))
+  layer = Chain(Chain(conv_bn((3, 3), 3, 32; stride = 2)...,
+                      conv_bn((3, 3), 32, 32)...,
+                      conv_bn((3, 3), 32, 64; pad = 1)...,
+                      MaxPool((3, 3), stride = 2),
+                      conv_bn((1, 1), 64, 80)...,
+                      conv_bn((3, 3), 80, 192)...,
+                      MaxPool((3, 3), stride = 2),
+                      inception_a(192, 32),
+                      inception_a(256, 64),
+                      inception_a(288, 64),
+                      inception_b(288),
+                      inception_c(768, 128),
+                      inception_c(768, 160),
+                      inception_c(768, 160),
+                      inception_c(768, 192),
+                      inception_d(768),
+                      inception_e(1280),
+                      inception_e(2048),
+                      AdaptiveMeanPool((1, 1))),
+                Chain(Dropout(0.2),
+                      flatten,
+                      Dense(2048, nclasses)))
 
   return layer
 end

src/resnet.jl

@@ -104,7 +104,7 @@ function resnet(; block, shortcut_config, channel_config, block_config, nclasses
                                 skip_projection(inplanes, outplanes[end], i != 1)))
     end
     inplanes = outplanes[end]
-    for j in 2:nrepeats
+    for _ in 2:nrepeats
       if shortcut_config == :A || shortcut_config == :B
         push!(layers, Parallel(+, block(inplanes, outplanes, false),
                                   skip_identity(inplanes, outplanes[end])))
@@ -116,11 +116,9 @@ function resnet(; block, shortcut_config, channel_config, block_config, nclasses
     end
     baseplanes *= 2
   end
-  push!(layers, AdaptiveMeanPool((1, 1)))
-  push!(layers, flatten)
-  push!(layers, Dense(inplanes, nclasses))
 
-  return Chain(layers...)
+  return Chain(Chain(layers..., AdaptiveMeanPool((1, 1))),
+               Chain(flatten, Dense(inplanes, nclasses)))
 end
 
 const resnet_config =

src/squeezenet.jl

@@ -28,21 +28,21 @@ Create a SqueezeNet
 ([reference](https://arxiv.org/abs/1602.07360v4)).
 """
 function squeezenet()
-  layers = Chain(Conv((3, 3), 3 => 64, relu, stride = 2),
-                 MaxPool((3, 3), stride = 2),
-                 fire(64, 16, 64, 64),
-                 fire(128, 16, 64, 64),
-                 MaxPool((3, 3), stride = 2),
-                 fire(128, 32, 128, 128),
-                 fire(256, 32, 128, 128),
-                 MaxPool((3, 3), stride = 2),
-                 fire(256, 48, 192, 192),
-                 fire(384, 48, 192, 192),
-                 fire(384, 64, 256, 256),
-                 fire(512, 64, 256, 256),
-                 Dropout(0.5),
-                 Conv((1, 1), 512 => 1000, relu),
-                 AdaptiveMeanPool((1, 1)),
+  layers = Chain(Chain(Conv((3, 3), 3 => 64, relu, stride = 2),
+                       MaxPool((3, 3), stride = 2),
+                       fire(64, 16, 64, 64),
+                       fire(128, 16, 64, 64),
+                       MaxPool((3, 3), stride = 2),
+                       fire(128, 32, 128, 128),
+                       fire(256, 32, 128, 128),
+                       MaxPool((3, 3), stride = 2),
+                       fire(256, 48, 192, 192),
+                       fire(384, 48, 192, 192),
+                       fire(384, 64, 256, 256),
+                       fire(512, 64, 256, 256),
+                       Dropout(0.5),
+                       Conv((1, 1), 512 => 1000, relu),
+                       AdaptiveMeanPool((1, 1))),
                  flatten)
 
   return layers

src/vgg.jl

@@ -94,7 +94,7 @@ function vgg(imsize; config, inchannels, batchnorm = false, nclasses, fcsize, dr
   conv = vgg_convolutional_layers(config, batchnorm, inchannels)
   imsize = outputsize(conv, (imsize..., inchannels); padbatch = true)[1:3]
   class = vgg_classifier_layers(imsize, nclasses, fcsize, dropout)
-  return Chain(conv..., class...)
+  return Chain(Chain(conv...), Chain(class...))
 end
 
 const vgg_config = Dict(:A => [(64,1), (128,1), (256,2), (512,2), (512,2)],