Adding some Untrained Models

src/Metalhead.jl

@@ -5,7 +5,12 @@ using Flux, Images, ImageFiltering, BSON, REPL, Requires, Statistics
 using Flux: @treelike
 
 # Models
-export VGG19, SqueezeNet, DenseNet, ResNet, GoogleNet
+export VGG19, VGG16, VGG13, VGG11, SqueezeNet, DenseNet121,
+       DenseNet169, DenseNet201, DenseNet264, ResNet18,
+       ResNet34, ResNet50, ResNet101, ResNet152, GoogleNet
+
+# Trained Models Loader
+export trained
 
 # Useful re-export from Images
 export load
@@ -30,7 +35,7 @@ include("display/terminal.jl")
 include("datasets/imagenet.jl")
 include("datasets/cifar10.jl")
 include("datasets/autodetect.jl")
-include("vgg19.jl")
+include("vgg.jl")
 include("squeezenet.jl")
 include("densenet.jl")
 include("resnet.jl")

src/densenet.jl

@@ -14,7 +14,7 @@ Bottleneck(in_planes, growth_rate) = Bottleneck(
 
 Transition(chs::Pair{<:Int, <:Int}) = Chain(BatchNorm(chs[1], relu),
                                             Conv((1, 1), chs),
-                                            x -> meanpool(x, (2, 2)))
+                                            MeanPool((2, 2)))
 
 function _make_dense_layers(block, in_planes, growth_rate, nblock)
   local layers = []
@@ -25,37 +25,13 @@ function _make_dense_layers(block, in_planes, growth_rate, nblock)
   Chain(layers...)
 end
 
-function _densenet(nblocks = [6, 12, 24, 16]; block = Bottleneck, growth_rate = 32, reduction = 0.5, num_classes = 1000)
-  num_planes = 2growth_rate
-  layers = []
-  push!(layers, Conv((7, 7), 3=>num_planes, stride = (2, 2), pad = (3, 3)))
-  push!(layers, BatchNorm(num_planes, relu))
-  push!(layers, x -> maxpool(x, (3, 3), stride = (2, 2), pad = (1, 1)))
-
-  for i in 1:3
-    push!(layers, _make_dense_layers(block, num_planes, growth_rate, nblocks[i]))
-    num_planes += nblocks[i] * growth_rate
-    out_planes = Int(floor(num_planes * reduction))
-    push!(layers, Transition(num_planes=>out_planes))
-    num_planes = out_planes
-  end
-
-  push!(layers, _make_dense_layers(block, num_planes, growth_rate, nblocks[4]))
-  num_planes += nblocks[4] * growth_rate
-  push!(layers, BatchNorm(num_planes, relu))
-
-  Chain(layers..., x -> meanpool(x, (7, 7)),
-        x -> reshape(x, :, size(x, 4)),
-        Dense(num_planes, num_classes), softmax)
-end
-
-function densenet_layers()
+function trained_densenet121_layers()
   weight = Metalhead.weights("densenet.bson")
   weights = Dict{Any, Any}()
   for ele in keys(weight)
     weights[string(ele)] = convert(Array{Float64, N} where N ,weight[ele])
   end
-  ls = _densenet()
+  ls = load_densenet(densenet_configs["densenet121"]...)
   ls[1].weight.data .= weights["conv1_w_0"][end:-1:1,:,:,:][:,end:-1:1,:,:]
   ls[2].β.data .= weights["conv1/bn_b_0"]
   ls[2].γ.data .= weights["conv1/bn_w_0"]
@@ -81,14 +57,88 @@ function densenet_layers()
   return ls
 end
 
-struct DenseNet <: ClassificationModel{ImageNet.ImageNet1k}
+function load_densenet(block, nblocks; growth_rate = 32, reduction = 0.5, num_classes = 1000)
+  num_planes = 2growth_rate
+  layers = []
+  push!(layers, Conv((7, 7), 3=>num_planes, stride = (2, 2), pad = (3, 3)))
+  push!(layers, BatchNorm(num_planes, relu))
+  push!(layers, MaxPool((3, 3), stride = (2, 2), pad = (1, 1)))
+
+  for i in 1:3
+    push!(layers, _make_dense_layers(block, num_planes, growth_rate, nblocks[i]))
+    num_planes += nblocks[i] * growth_rate
+    out_planes = Int(floor(num_planes * reduction))
+    push!(layers, Transition(num_planes=>out_planes))
+    num_planes = out_planes
+  end
+
+  push!(layers, _make_dense_layers(block, num_planes, growth_rate, nblocks[4]))
+  num_planes += nblocks[4] * growth_rate
+  push!(layers, BatchNorm(num_planes, relu))
+
+  Chain(layers..., x -> MeanPool((7, 7)),
+        x -> reshape(x, :, size(x, 4)),
+        Dense(num_planes, num_classes), softmax)
+end
+
+densenet_configs =
+  Dict("densenet121" => (Bottleneck, [6, 12, 24, 16]),
+       "densenet169" => (Bottleneck, [6, 12, 32, 32]),
+       "densenet201" => (Bottleneck, [6, 12, 48, 32]),
+       "densenet264" => (Bottleneck, [6, 12, 64, 48]))
+
+struct DenseNet121 <: ClassificationModel{ImageNet.ImageNet1k}
+  layers::Chain
+end
+
+DenseNet121() = DenseNet121(load_densenet(densenet_configs["densenet121"]...))
+
+trained(::Type{DenseNet121}) = DenseNet121(trained_densenet121_layers())
+
+Base.show(io::IO, ::DenseNet121) = print(io, "DenseNet264()")
+
+@treelike DenseNet121
+
+(m::DenseNet121)(x) = m.layers(x)
+
+struct DenseNet169 <: ClassificationModel{ImageNet.ImageNet1k}
+  layers::Chain
+end
+
+DenseNet169() = DenseNet169(load_densenet(densenet_configs["densenet169"]...))
+
+trained(::Type{DenseNet169}) = error("Pretrained Weights for DenseNet169 are not available")
+
+Base.show(io::IO, ::DenseNet169) = print(io, "DenseNet169()")
+
+@treelike DenseNet169
+
+(m::DenseNet169)(x) = m.layers(x)
+
+struct DenseNet201 <: ClassificationModel{ImageNet.ImageNet1k}
   layers::Chain
 end
 
-DenseNet() = DenseNet(densenet_layers())
+DenseNet201() = DenseNet201(load_densenet(densenet_configs["densenet201"]...))
+
+trained(::Type{DenseNet201}) = error("Pretrained Weights for DenseNet201 are not available")
+
+Base.show(io::IO, ::DenseNet201) = print(io, "DenseNet201()")
+
+@treelike DenseNet201
+
+(m::DenseNet201)(x) = m.layers(x)
+
+struct DenseNet264 <: ClassificationModel{ImageNet.ImageNet1k}
+  layers::Chain
+end
+
+DenseNet264() = DenseNet264(load_densenet(densenet_configs["densenet264"]..., growth_rate=48))
+
+trained(::Type{DenseNet264}) = error("Pretrained Weights for DenseNet264 are not available")
 
-Base.show(io::IO, ::DenseNet) = print(io, "DenseNet()")
+Base.show(io::IO, ::DenseNet264) = print(io, "DenseNet264()")
 
-@treelike DenseNet
+@treelike DenseNet264
 
-(m::DenseNet)(x) = m.layers(x)
+(m::DenseNet264)(x) = m.layers(x)

src/resnet.jl

@@ -16,9 +16,8 @@ function ResidualBlock(filters, kernels::Array{Tuple{Int,Int}}, pads::Array{Tupl
   ResidualBlock(Tuple(conv_layers),Tuple(norm_layers),shortcut)
 end
 
-function ResidualBlock(filters, kernels::Array{Int}, pads::Array{Int}, strides::Array{Int}, shortcut = identity)
+ResidualBlock(filters, kernels::Array{Int}, pads::Array{Int}, strides::Array{Int}, shortcut = identity) =
   ResidualBlock(filters, [(i,i) for i in kernels], [(i,i) for i in pads], [(i,i) for i in strides], shortcut)
-end
 
 function (block::ResidualBlock)(input)
   local value = copy.(input)
@@ -28,55 +27,39 @@ function (block::ResidualBlock)(input)
   relu.(((block.norm_layers[end])((block.conv_layers[end])(value))) + block.shortcut(input))
 end
 
-function Bottleneck(filters::Int, downsample::Bool = false, res_top::Bool = false)
-  if(!downsample && !res_top)
-    return ResidualBlock([4 * filters, filters, filters, 4 * filters], [1,3,1], [0,1,0], [1,1,1])
-  elseif(downsample && res_top)
-    return ResidualBlock([filters, filters, filters, 4 * filters], [1,3,1], [0,1,0], [1,1,1], Chain(Conv((1,1), filters=>4 * filters, pad = (0,0), stride = (1,1)), BatchNorm(4 * filters)))
-  else
-    shortcut = Chain(Conv((1,1), 2 * filters=>4 * filters, pad = (0,0), stride = (2,2)), BatchNorm(4 * filters))
-    return ResidualBlock([2 * filters, filters, filters, 4 * filters], [1,3,1], [0,1,0], [1,1,2], shortcut)
+function BasicBlock(filters::Int, downsample::Bool = false, res_top::Bool = false)
+  if !downsample || res_top
+    return ResidualBlock([filters for i in 1:3], [3,3], [1,1], [1,1])
   end
+  shortcut = Chain(Conv((3,3), filters÷2=>filters, pad = (1,1), stride = (2,2)), BatchNorm(filters))
+  ResidualBlock([filters÷2, filters, filters], [3,3], [1,1], [1,2], shortcut)
 end
 
-function resnet50()
-  local layers = [3, 4, 6, 3]
-  local layer_arr = []
-
-  push!(layer_arr, Conv((7,7), 3=>64, pad = (3,3), stride = (2,2)))
-  push!(layer_arr, x -> maxpool(x, (3,3), pad = (1,1), stride = (2,2)))
-
-  initial_filters = 64
-  for i in 1:length(layers)
-    push!(layer_arr, Bottleneck(initial_filters, true, i==1))
-    for j in 2:layers[i]
-      push!(layer_arr, Bottleneck(initial_filters))
-    end
-    initial_filters *= 2
+function Bottleneck(filters::Int, downsample::Bool = false, res_top::Bool = false)
+  if !downsample && !res_top
+    ResidualBlock([4 * filters, filters, filters, 4 * filters], [1,3,1], [0,1,0], [1,1,1])
+  elseif downsample && res_top
+    ResidualBlock([filters, filters, filters, 4 * filters], [1,3,1], [0,1,0], [1,1,1], Chain(Conv((1,1), filters=>4 * filters, pad = (0,0), stride = (1,1)), BatchNorm(4 * filters)))
+  else
+    shortcut = Chain(Conv((1,1), 2 * filters=>4 * filters, pad = (0,0), stride = (2,2)), BatchNorm(4 * filters))
+    ResidualBlock([2 * filters, filters, filters, 4 * filters], [1,3,1], [0,1,0], [1,1,2], shortcut)
   end
-
-  push!(layer_arr, x -> meanpool(x, (7,7)))
-  push!(layer_arr, x -> reshape(x, :, size(x,4)))
-  push!(layer_arr, (Dense(2048, 1000)))
-  push!(layer_arr, softmax)
-
-  Chain(layer_arr...)
 end
 
-function resnet_layers()
+function trained_resnet50_layers()
   weight = Metalhead.weights("resnet.bson")
   weights = Dict{Any ,Any}()
   for ele in keys(weight)
     weights[string(ele)] = convert(Array{Float64, N} where N, weight[ele])
   end
-  ls = resnet50()
-  ls[1].weight.data .= weights["gpu_0/conv1_w_0"][end:-1:1,:,:,:][:,end:-1:1,:,:]
+  ls = load_resnet(resnet_configs["resnet50"]...)
+  ls[1].weight.data .= flipkernel(weights["gpu_0/conv1_w_0"])
   count = 2
   for j in [3:5, 6:9, 10:15, 16:18]
     for p in j
-      ls[p].conv_layers[1].weight.data .= weights["gpu_0/res$(count)_$(p-j[1])_branch2a_w_0"][end:-1:1,:,:,:][:,end:-1:1,:,:]
-      ls[p].conv_layers[2].weight.data .= weights["gpu_0/res$(count)_$(p-j[1])_branch2b_w_0"][end:-1:1,:,:,:][:,end:-1:1,:,:]
-      ls[p].conv_layers[3].weight.data .= weights["gpu_0/res$(count)_$(p-j[1])_branch2c_w_0"][end:-1:1,:,:,:][:,end:-1:1,:,:]
+      ls[p].conv_layers[1].weight.data .= flipkernel(weights["gpu_0/res$(count)_$(p-j[1])_branch2a_w_0"])
+      ls[p].conv_layers[2].weight.data .= flipkernel(weights["gpu_0/res$(count)_$(p-j[1])_branch2b_w_0"])
+      ls[p].conv_layers[3].weight.data .= flipkernel(weights["gpu_0/res$(count)_$(p-j[1])_branch2c_w_0"])
     end
     count += 1
   end
@@ -85,14 +68,107 @@ function resnet_layers()
   return ls
 end
 
-struct ResNet <: ClassificationModel{ImageNet.ImageNet1k}
+function load_resnet(Block, layers, initial_filters::Int = 64, nclasses::Int = 1000)
+  local top = []
+  local residual = []
+  local bottom = []
+
+  push!(top, Conv((7,7), 3=>initial_filters, pad = (3,3), stride = (2,2)))
+  push!(top, MaxPool((3,3), pad = (1,1), stride = (2,2)))
+
+  for i in 1:length(layers)
+    push!(residual, Block(initial_filters, true, i==1))
+    for j in 2:layers[i]
+      push!(residual, Block(initial_filters))
+    end
+    initial_filters *= 2
+  end
+
+  push!(bottom, MeanPool((7,7)))
+  push!(bottom, x -> reshape(x, :, size(x,4)))
+  if Block == Bottleneck
+    push!(bottom, (Dense(2048, nclasses)))
+  else
+    push!(bottom, (Dense(512, nclasses)))
+  end
+  push!(bottom, softmax)
+
+  Chain(top..., residual..., bottom...)
+end
+
+resnet_configs =
+  Dict("resnet18" => (BasicBlock, [2, 2, 2, 2]),
+       "resnet34" => (BasicBlock, [3, 4, 6, 3]),
+       "resnet50" => (Bottleneck, [3, 4, 6, 3]),
+       "resnet101" => (Bottleneck, [3, 4, 23, 3]),
+       "resnet152" => (Bottleneck, [3, 8, 36, 3]))
+
+struct ResNet18 <: ClassificationModel{ImageNet.ImageNet1k}
+  layers::Chain
+end
+
+ResNet18() = ResNet18(load_resnet(resnet_configs["resnet18"]...))
+
+trained(::Type{ResNet18}) = error("Pretrained Weights for ResNet18 are not available")
+
+Base.show(io::IO, ::ResNet18) = print(io, "ResNet18()")
+
+@treelike ResNet18
+
+(m::ResNet18)(x) = m.layers(x)
+
+struct ResNet34 <: ClassificationModel{ImageNet.ImageNet1k}
+  layers::Chain
+end
+
+ResNet34() = ResNet34(load_resnet(resnet_configs["resnet34"]...))
+
+trained(::Type{ResNet34}) = error("Pretrained Weights for ResNet34 are not available")
+
+Base.show(io::IO, ::ResNet34) = print(io, "ResNet34()")
+
+@treelike ResNet34
+
+(m::ResNet34)(x) = m.layers(x)
+
+struct ResNet50 <: ClassificationModel{ImageNet.ImageNet1k}
   layers::Chain
 end
 
-ResNet() = ResNet(resnet_layers())
+ResNet50() = ResNet50(load_resnet(resnet_configs["resnet50"]...))
+
+trained(::Type{ResNet50}) = ResNet50(trained_resnet50_layers())
+
+Base.show(io::IO, ::ResNet50) = print(io, "ResNet50()")
+
+@treelike ResNet50
+
+(m::ResNet50)(x) = m.layers(x)
+
+struct ResNet101 <: ClassificationModel{ImageNet.ImageNet1k}
+  layers::Chain
+end
+
+ResNet101() = ResNet101(load_resnet(resnet_configs["resnet101"]...))
+
+trained(::Type{ResNet101}) = error("Pretrained Weights for ResNet101 are not available")
+
+Base.show(io::IO, ::ResNet101) = print(io, "ResNet101()")
+
+@treelike ResNet101
+
+(m::ResNet101)(x) = m.layers(x)
+
+struct ResNet152 <: ClassificationModel{ImageNet.ImageNet1k}
+  layers::Chain
+end
+
+ResNet152() = ResNet152(load_resnet(resnet_configs["resnet152"]...))
+
+trained(::Type{ResNet152}) = error("Pretrained Weights for ResNet152 are not available")
 
-Base.show(io::IO, ::ResNet) = print(io, "ResNet()")
+Base.show(io::IO, ::ResNet152) = print(io, "ResNet152()")
 
-@treelike ResNet
+@treelike ResNet152
 
-(m::ResNet)(x) = m.layers(x)
+(m::ResNet152)(x) = m.layers(x)

src/utils.jl

@@ -21,6 +21,9 @@ function weights(name)
   BSON.load(joinpath(deps, name))
 end
 
+# TODO: Remove after NNlib supports flip kernel through https://github.com/FluxML/NNlib.jl/pull/53
+flipkernel(x::AbstractArray) = x[end:-1:1, end:-1:1, :, :]
+
 load_img(im::AbstractMatrix{<:Color}) = im
 load_img(str::AbstractString) = load(str)
 load_img(val::ValidationImage) = load_img(val.img)