Feature registries for models

Project.toml

@@ -5,6 +5,7 @@ version = "0.7.0"
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
+FeatureRegistries = "c6aefb4f-3ac3-4095-8805-528476b02c02"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
@@ -19,8 +20,9 @@ Flux = "0.13"
 Functors = "0.2"
 MLUtils = "0.2"
 NNlib = "0.7.34, 0.8"
-julia = "1.6"
 NeuralAttentionlib = "0.0"
+FeatureRegistries = "0.1"
+julia = "1.6"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/Metalhead.jl

@@ -1,8 +1,8 @@
 module Metalhead
 
 using Flux
-using Flux: outputsize, Zygote
 using Functors
+using FeatureRegistries
 using BSON
 using Artifacts, LazyArtifacts
 using Statistics
@@ -36,24 +36,22 @@ include("other/mlpmixer.jl")
 # ViT-based models
 include("vit-based/vit.jl")
 
-include("pretrain.jl")
+const _MODEL_NAMES = [:AlexNet, :VGG, :ResNet, :GoogLeNet, :Inception3, :SqueezeNet, :DenseNet,
+                      :ResNeXt, :MobileNetv1, :MobileNetv2, :MobileNetv3, :MLPMixer, :ResMLP,
+                      :gMLP, :ViT, :ConvNeXt, :ConvMixer]
 
-export  AlexNet,
-        VGG, VGG11, VGG13, VGG16, VGG19,
-        ResNet, ResNet18, ResNet34, ResNet50, ResNet101, ResNet152,
-        GoogLeNet, Inception3, SqueezeNet,
-        DenseNet, DenseNet121, DenseNet161, DenseNet169, DenseNet201,
-        ResNeXt,
-        MobileNetv1, MobileNetv2, MobileNetv3,
-        MLPMixer, ResMLP, gMLP,
-        ViT,
-        ConvNeXt, ConvMixer
-
-# use Flux._big_show to pretty print large models
-for T in (:AlexNet, :VGG, :ResNet, :GoogLeNet, :Inception3, :SqueezeNet, :DenseNet, :ResNeXt, 
-          :MobileNetv1, :MobileNetv2, :MobileNetv3,
-          :MLPMixer, :ResMLP, :gMLP, :ViT, :ConvNeXt, :ConvMixer)
+# Export models, use Flux._big_show to pretty print large models
+for T in _MODEL_NAMES
+  @eval export $T
   @eval Base.show(io::IO, ::MIME"text/plain", model::$T) = _maybe_big_show(io, model)
 end
 
+# Export deprecated models
+export VGG11, VGG13, VGG16, VGG19,
+       ResNet18, ResNet34, ResNet50, ResNet101, ResNet152,
+       DenseNet121, DenseNet161, DenseNet169, DenseNet201
+
+include("pretrain.jl")
+include("registries.jl")
+
 end # module

src/convnets/inception.jl

@@ -145,9 +145,6 @@ Create an Inception-v3 model ([reference](https://arxiv.org/abs/1512.00567v3)).
 
 # Arguments
 - `nclasses`: the number of output classes
-
-!!! warning
-    `inception3` does not currently support pretrained weights.
 """
 function inception3(; nclasses = 1000)
   layer = Chain(Chain(conv_bn((3, 3), 3, 32; stride = 2),

src/convnets/mobilenet.jl

@@ -67,6 +67,9 @@ Create a MobileNetv1 model with the baseline configuration
 ([reference](https://arxiv.org/abs/1704.04861v1)).
 Set `pretrain` to `true` to load the pretrained weights for ImageNet.
 
+!!! warning
+    `MobileNetv1` does not currently support pretrained weights.
+
 # Arguments
 - `width_mult`: Controls the number of output feature maps in each block
                 (with 1.0 being the default in the paper;
@@ -160,7 +163,10 @@ end
 
 Create a MobileNetv2 model with the specified configuration.
 ([reference](https://arxiv.org/abs/1801.04381)).
-Set `pretrain` to `true` to load the pretrained weights for ImageNet. 
+Set `pretrain` to `true` to load the pretrained weights for ImageNet.
+
+!!! warning
+    `MobileNetv2` does not currently support pretrained weights.
 
 # Arguments
 - `width_mult`: Controls the number of output feature maps in each block
@@ -282,6 +288,9 @@ Create a MobileNetv3 model with the specified configuration.
 ([reference](https://arxiv.org/abs/1905.02244)).
 Set `pretrain = true` to load the model with pre-trained weights for ImageNet.
 
+!!! warning
+    `MobileNetv3` does not currently support pretrained weights.
+
 # Arguments
 - `mode`: :small or :large for the size of the model (see paper).
 - `width_mult`: Controls the number of output feature maps in each block

src/convnets/resnext.jl

@@ -106,7 +106,7 @@ Create a ResNeXt model with specified configuration. Currently supported values
 Set `pretrain = true` to load the model with pre-trained weights for ImageNet.
 
 !!! warning
-  `ResNeXt` does not currently support pretrained weights.
+    `ResNeXt` does not currently support pretrained weights.
 
 See also [`Metalhead.resnext`](#).
 """

src/convnets/vgg.jl

@@ -89,7 +89,7 @@ Create a VGG model
 """
 function vgg(imsize; config, inchannels, batchnorm = false, nclasses, fcsize, dropout)
   conv = vgg_convolutional_layers(config, batchnorm, inchannels)
-  imsize = outputsize(conv, (imsize..., inchannels); padbatch = true)[1:3]
+  imsize = Flux.outputsize(conv, (imsize..., inchannels); padbatch = true)[1:3]
   class = vgg_classifier_layers(imsize, nclasses, fcsize, dropout)
   return Chain(Chain(conv), class)
 end

src/layers/Layers.jl

@@ -1,7 +1,6 @@
 module Layers
 
 using Flux
-using Flux: outputsize, Zygote
 using Functors
 using Statistics
 using MLUtils

src/registries.jl

@@ -0,0 +1,27 @@
+using FeatureRegistries: Registry, Field
+
+commafiy(n::Integer) =
+         join(reverse(join.(reverse.(Iterators.partition(digits(n), 3)))), ',')
+
+function model_registry()
+  registry = Registry((;
+        serial_number = Field(Integer, name = "Serial Number", description = "Serial number"),
+        model_name = Field(
+            String,
+            name = "Model name",
+            description = "The name of the model",
+        ),
+        parameters = Field(
+            Any,
+            name = "Parameters",
+            description = "The number of parameters in the model",
+        )),
+        name = "Models", id = :model_name,)
+
+  for (i, name) in enumerate(_MODEL_NAMES)
+    param_string = Base.Docs.doc(Base.Docs.Binding(Main, name))
+    push!(registry, (serial_number = i, model_name = String(name), parameters = param_string))
+  end
+
+  return registry
+end