Merge pull request #154 from theabhirath/densenet-classes

Expose `nclasses` option for `DenseNet`

Author: darsnack

Project.toml

@@ -1,6 +1,6 @@
 name = "Metalhead"
 uuid = "dbeba491-748d-5e0e-a39e-b530a07fa0cc"
-version = "0.7.0"
+version = "0.7.1"
 
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

src/convnets/densenet.jl

@@ -147,9 +147,9 @@ Set `pretrain = true` to load the model with pre-trained weights for ImageNet.
 
 See also [`Metalhead.densenet`](#).
 """
-function DenseNet(config::Int = 121; pretrain = false)
+function DenseNet(config::Int = 121; pretrain = false, nclasses = 1000)
   @assert config in keys(densenet_config) "`config` must be one out of $(sort(collect(keys(densenet_config))))."
-  model = DenseNet(densenet_config[config])
+  model = DenseNet(densenet_config[config]; nclasses = nclasses)
 
   pretrain && loadpretrain!(model, string("DenseNet", config))
   return model