Minor formatting tweaks

Author: theabhirath

src/Metalhead.jl

@@ -37,7 +37,7 @@ export  AlexNet,
 
 # use Flux._big_show to pretty print large models
 for T in (:AlexNet, :VGG, :ResNet, :GoogLeNet, :Inception3, :SqueezeNet, :DenseNet, :ResNeXt, 
-    :MobileNetv2, :MobileNetv3, :MLPMixer)
+          :MobileNetv2, :MobileNetv3, :MLPMixer)
 @eval Base.show(io::IO, ::MIME"text/plain", model::$T) = _maybe_big_show(io, model)
 end
 

src/other/mlpmixer.jl

@@ -47,7 +47,7 @@ function mlpmixer(imsize::NTuple{2} = (256, 256); inchannels = 3, patch_size = 1
 
   classification_head = Chain(_seconddimmean, Dense(planes, nclasses))
 
-  return Chain(layers..., classification_head)
+  return Chain(Chain(layers...), classification_head)
 end
 
 struct MLPMixer
@@ -56,7 +56,7 @@ end
 
 """
     MLPMixer(imsize::NTuple{2} = (256, 256); inchannels = 3, patch_size = 16, planes = 512, 
-             depth = 12, expansion_factor = 4, dropout = 0., pretrain = false, nclasses = 1000)
+             depth = 12, expansion_factor = 4, dropout = 0., nclasses = 1000)
 
 Creates a model with the MLPMixer architecture.
 ([reference](https://arxiv.org/pdf/2105.01601)).
@@ -69,18 +69,19 @@ Creates a model with the MLPMixer architecture.
 - depth: the number of blocks in the main model
 - expansion_factor: the number of channels in each block
 - dropout: the dropout rate
-- pretrain: whether to load the pre-trained weights for ImageNet
 - nclasses: the number of classes in the output
 """
 function MLPMixer(imsize::NTuple{2} = (256, 256); inchannels = 3, patch_size = 16, planes = 512, 
-                  depth = 12, expansion_factor = 4, dropout = 0., pretrain = false, nclasses = 1000)
+                  depth = 12, expansion_factor = 4, dropout = 0., nclasses = 1000)
 
   layers = mlpmixer(imsize; inchannels, patch_size, planes, depth, expansion_factor, dropout, 
                     nclasses)
-  pretrain && loadpretrain!(layers, string("MLPMixer"))
   MLPMixer(layers)
 end
 
+@functor MLPMixer
+
 (m::MLPMixer)(x) = m.layers(x)
 
-@functor MLPMixer
+backbone(m::MLPMixer) = m.layers[1]
+classifier(m::MLPMixer) = m.layers[2:end]

test/other.jl

@@ -2,6 +2,6 @@ using Metalhead, Test
 using Flux
 
 @testset "MLPMixer" begin
-    @test size(MLPMixer()(rand(Float32, 256, 256, 3, 67))) == (1000, 67)
-    @test_skip gradtest(MLPMixer(), rand(Float32, 256, 256, 3, 67))
+    @test size(MLPMixer()(rand(Float32, 256, 256, 3, 2))) == (1000, 2)
+    @test_skip gradtest(MLPMixer(), rand(Float32, 256, 256, 3, 2))
 end