Merge branch 'refine' of https://github.com/theabhirath/Metalhead.jl into refine

Author: theabhirath

src/layers/mlp.jl

@@ -47,7 +47,7 @@ end
 gated_mlp_block(::typeof(identity), args...; kwargs...) = mlp_block(args...; kwargs...)
 
 """
-    create_classifier(inplanes::Integer, nclasses::Integer, activation = relu;
+    create_classifier(inplanes::Integer, nclasses::Integer, activation = identity;
                       pool_layer = AdaptiveMeanPool((1, 1)),
                       dropout_rate = 0.0, use_conv::Bool = false)
 
@@ -64,7 +64,7 @@ Creates a classifier head to be used for models.
   - `use_conv`: whether to use a 1x1 convolutional layer instead of a `Dense` layer.
 """
 function create_classifier(inplanes::Integer, nclasses::Integer, activation = identity;
-                           use_conv::Bool = falsepool_layer = AdaptiveMeanPool((1, 1)),
+                           use_conv::Bool = false, pool_layer = AdaptiveMeanPool((1, 1)),
                            dropout_rate = nothing)
     # Pooling
     flatten_in_pool = !use_conv && pool_layer !== identity

src/mixers/mlpmixer.jl

@@ -34,8 +34,8 @@ function mixerblock(planes::Integer, npatches::Integer; mlp_layer = mlp_block,
 end
 
 """
-MLPMixer(size::Symbol; patch_size::Dims{2} = (16, 16), imsize::Dims{2} = (224, 224),
-inchannels::Integer = 3, nclasses::Integer = 1000)
+    MLPMixer(size::Symbol; patch_size::Dims{2} = (16, 16), imsize::Dims{2} = (224, 224),
+             inchannels::Integer = 3, nclasses::Integer = 1000)
 
 Creates a model with the MLPMixer architecture.
 ([reference](https://arxiv.org/pdf/2105.01601)).

src/vit-based/vit.jl

@@ -1,5 +1,5 @@
 """
-transformer_encoder(planes, depth, nheads; mlp_ratio = 4.0, dropout_rate = 0.)
+    transformer_encoder(planes, depth, nheads; mlp_ratio = 4.0, dropout_rate = 0.)
 
 Transformer as used in the base ViT architecture.
 ([reference](https://arxiv.org/abs/2010.11929)).
@@ -99,12 +99,10 @@ struct ViT
 end
 @functor ViT
 
-function ViT(mode::Symbol = :base; imsize::Dims{2} = (256, 256),
-             patch_size::Dims{2} = (16, 16),
+function ViT(mode::Symbol; imsize::Dims{2} = (256, 256), patch_size::Dims{2} = (16, 16),
              inchannels::Integer = 3, nclasses::Integer = 1000)
     _checkconfig(mode, keys(VIT_CONFIGS))
-    kwargs = VIT_CONFIGS[mode]
-    layers = vit(imsize; inchannels, patch_size, nclasses, kwargs...)
+    layers = vit(imsize; inchannels, patch_size, nclasses, VIT_CONFIGS[mode]...)
     return ViT(layers)
 end