Merge pull request #103 from theabhirath/mlpmixer

Implementation of MLPMixer

Author: darsnack

Project.toml

@@ -8,17 +8,18 @@ BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 BSON = "0.3.2"
 Flux = "0.12"
 Functors = "0.2"
-julia = "1.4"
 NNlib = "0.7.34"
+julia = "1.4"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 
 [publish]
 title = "Metalhead.jl"

README.md

@@ -41,7 +41,7 @@
 | ResNeXt-152                                      | [`ResNeXt`](https://fluxml.ai/Metalhead.jl/dev/docstrings/Metalhead.ResNeXt.html)          | N            |
 | [MobileNetv2](https://arxiv.org/abs/1801.04381)  | [`MobileNetv2`](https://fluxml.ai/Metalhead.jl/dev/docstrings/Metalhead.MobileNetv2.html)  | N            |
 | [MobileNetv3](https://arxiv.org/abs/1905.02244)  | [`MobileNetv3`](https://fluxml.ai/Metalhead.jl/dev/docstrings/Metalhead.MobileNetv3.html)  | N            |
-
+| [MLPMixer](https://arxiv.org/pdf/2105.01601)     | [`MLPMixer`](https://fluxml.ai/Metalhead.jl/dev/docstrings/Metalhead.MLPMixer.html)     | N            |
 
 ## Getting Started
 

src/Metalhead.jl

@@ -5,32 +5,39 @@ using Flux: outputsize, Zygote
 using Functors
 using BSON
 using Artifacts, LazyArtifacts
+using Statistics
 
 import Functors
 
-# Models
 include("utilities.jl")
-include("alexnet.jl")
-include("vgg.jl")
-include("resnet.jl")
-include("googlenet.jl")
-include("inception.jl")
-include("squeezenet.jl")
-include("densenet.jl")
-include("resnext.jl")
-include("mobilenet.jl")
+include("layers.jl")
+
+# CNN models
+include("convnets/alexnet.jl")
+include("convnets/vgg.jl")
+include("convnets/inception.jl")
+include("convnets/googlenet.jl")
+include("convnets/resnet.jl")
+include("convnets/resnext.jl")
+include("convnets/densenet.jl")
+include("convnets/squeezenet.jl")
+include("convnets/mobilenet.jl")
+
+# Other models
+include("other/mlpmixer.jl")
 
 export  AlexNet,
         VGG, VGG11, VGG13, VGG16, VGG19,
         ResNet, ResNet18, ResNet34, ResNet50, ResNet101, ResNet152,
         GoogLeNet, Inception3, SqueezeNet,
         DenseNet, DenseNet121, DenseNet161, DenseNet169, DenseNet201,
         ResNeXt,
-        MobileNetv2, MobileNetv3
+        MobileNetv2, MobileNetv3,
+        MLPMixer
 
 # use Flux._big_show to pretty print large models
 for T in (:AlexNet, :VGG, :ResNet, :GoogLeNet, :Inception3, :SqueezeNet, :DenseNet, :ResNeXt, 
-          :MobileNetv2, :MobileNetv3)
+          :MobileNetv2, :MobileNetv3, :MLPMixer)
   @eval Base.show(io::IO, ::MIME"text/plain", model::$T) = _maybe_big_show(io, model)
 end