Actually rewrite VGG to use vgg

FluxML · Jan 4, 2024 · 09f91c0 · 09f91c0
1 parent ee2d8bf
commit 09f91c0
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/src/Metalhead.jl b/src/Metalhead.jl
@@ -72,6 +72,9 @@ include("vit-based/vit.jl")
 # Load pretrained weights
 include("pretrain.jl")
 
+# deprecated
+include("deprecations.jl")
+
 # export model functions
 export AlexNet, VGG, ResNet, WideResNet, ResNeXt, DenseNet,
        GoogLeNet, Inceptionv3, Inceptionv4, InceptionResNetv2, Xception,

diff --git a/src/convnets/vgg.jl b/src/convnets/vgg.jl
@@ -96,12 +96,10 @@ function vgg(imsize::Dims{2}; config, batchnorm::Bool = false, fcsize::Integer =
     return Chain(Chain(conv...), class)
 end
 
-const VGG_CONV_CONFIGS = Dict(:A => [(64, 1), (128, 1), (256, 2), (512, 2), (512, 2)],
-                              :B => [(64, 2), (128, 2), (256, 2), (512, 2), (512, 2)],
-                              :D => [(64, 2), (128, 2), (256, 3), (512, 3), (512, 3)],
-                              :E => [(64, 2), (128, 2), (256, 4), (512, 4), (512, 4)])
-
-const VGG_CONFIGS = Dict(11 => :A, 13 => :B, 16 => :D, 19 => :E)
+const VGG_CONFIGS = Dict(11 => [(64, 1), (128, 1), (256, 2), (512, 2), (512, 2)],
+                         13 => [(64, 2), (128, 2), (256, 2), (512, 2), (512, 2)],
+                         16 => [(64, 2), (128, 2), (256, 3), (512, 3), (512, 3)],
+                         19 => [(64, 2), (128, 2), (256, 4), (512, 4), (512, 4)])
 
 """
     VGG(depth::Integer; pretrain::Bool = false, batchnorm::Bool = false,
@@ -132,8 +130,7 @@ end
 function VGG(depth::Integer; pretrain::Bool = false, batchnorm::Bool = false,
              inchannels::Integer = 3, nclasses::Integer = 1000)
     _checkconfig(depth, keys(VGG_CONFIGS))
-    model = VGG((224, 224); config = VGG_CONV_CONFIGS[VGG_CONFIGS[depth]], batchnorm,
-                inchannels, nclasses)
+    model = vgg((224, 224); config = VGG_CONFIGS[depth], batchnorm, inchannels, nclasses)
     if pretrain
         artifact_name = string("vgg", depth)
         if batchnorm

diff --git a/src/deprecations.jl b/src/deprecations.jl
@@ -6,4 +6,4 @@ function VGG(imsize::Dims{2}; config, batchnorm::Bool = false, dropout_prob = 0.
                 inchannels, batchnorm, nclasses)` instead for the same functionality.", :VGG)
     layers = vgg(imsize; config, inchannels, batchnorm, nclasses, dropout_prob)
     return VGG(layers)
-end
+end