Adding UNet Model

.github/workflows/CI.yml

@@ -39,7 +39,7 @@ jobs:
           - 'r"/*/SEResNet*"'
           - '[r"Res2Net", r"Res2NeXt"]'
           - '"Inception"'
-          - '"DenseNet"'
+          - '["UNet", "DenseNet"]'
           - '["ConvNeXt", "ConvMixer"]'
           - 'r"Mixers"'
           - 'r"ViTs"'

docs/src/index.md

@@ -36,6 +36,7 @@ julia> ]add Metalhead
 | [ViT](https://arxiv.org/abs/2010.11929)                | [`ViT`](@ref)               | N            |
 | [ConvNeXt](https://arxiv.org/abs/2201.03545)           | [`ConvNeXt`](@ref)          | N            |
 | [ConvMixer](https://arxiv.org/abs/2201.09792)          | [`ConvMixer`](@ref)         | N            |
+| [UNet](https://arxiv.org/abs/1505.04597v1)             | [`UNet`](@ref)              | N            |
 
 To contribute new models, see our [contributing docs](@ref Contributing-to-Metalhead.jl).
 

src/Metalhead.jl

@@ -53,6 +53,7 @@ include("convnets/densenet.jl")
 include("convnets/squeezenet.jl")
 include("convnets/convnext.jl")
 include("convnets/convmixer.jl")
+include("convnets/unet.jl")
 
 # Mixers
 include("mixers/core.jl")
@@ -73,15 +74,15 @@ export AlexNet, VGG, VGG11, VGG13, VGG16, VGG19,
        GoogLeNet, Inception3, Inceptionv3, Inceptionv4, InceptionResNetv2, Xception,
        SqueezeNet, MobileNetv1, MobileNetv2, MobileNetv3, MNASNet,
        EfficientNet, EfficientNetv2, ConvMixer, ConvNeXt,
-       MLPMixer, ResMLP, gMLP, ViT
+       MLPMixer, ResMLP, gMLP, ViT, UNet
 
 # use Flux._big_show to pretty print large models
 for T in (:AlexNet, :VGG, :SqueezeNet, :ResNet, :WideResNet, :ResNeXt,
           :SEResNet, :SEResNeXt, :Res2Net, :Res2NeXt, :GoogLeNet, :DenseNet,
           :Inceptionv3, :Inceptionv4, :InceptionResNetv2, :Xception,
           :MobileNetv1, :MobileNetv2, :MobileNetv3, :MNASNet,
           :EfficientNet, :EfficientNetv2, :ConvMixer, :ConvNeXt,
-          :MLPMixer, :ResMLP, :gMLP, :ViT)
+          :MLPMixer, :ResMLP, :gMLP, :ViT, :UNet)
     @eval Base.show(io::IO, ::MIME"text/plain", model::$T) = _maybe_big_show(io, model)
 end
 

src/convnets/unet.jl

@@ -0,0 +1,123 @@
+function PixelShuffleICNR(inplanes, outplanes; r = 2)
+    return Chain(Chain(basic_conv_bn((1, 1), inplanes, outplanes * (r^2)...)),
+                 Flux.PixelShuffle(r))
+end
+
+function UNetCombineLayer(inplanes, outplanes)
+    return Chain(Chain(basic_conv_bn((3, 3), inplanes, outplanes; pad = 1)...),
+                 Chain(basic_conv_bn((3, 3), outplanes, outplanes; pad = 1)...))
+end
+
+function UNetMiddleBlock(inplanes)
+    return Chain(Chain(basic_conv_bn((3, 3), inplanes, 2 * inplanes; pad = 1)...),
+                 Chain(basic_conv_bn((3, 3), 2 * inplanes, inplanes; pad = 1)...))
+end
+
+function UNetFinalBlock(inplanes, outplanes)
+    return Chain(basicblock(inplanes, inplanes; reduction_factor = 1),
+                 Chain(basic_conv_bn((1, 1), inplanes, outplanes)...))
+end
+
+function unetlayers(layers, sz; outplanes = nothing, skip_upscale = 0,
+                    m_middle = _ -> (identity,))
+    isempty(layers) && return m_middle(sz[end - 1])
+
+    layer, layers = layers[1], layers[2:end]
+    outsz = Flux.outputsize(layer, sz)
+    does_downscale = sz[1] ÷ 2 == outsz[1]
+
+    if !does_downscale
+        return Chain(layer, unetlayers(layers, outsz; outplanes, skip_upscale)...)
+    elseif does_downscale && skip_upscale > 0
+        return Chain(layer,
+                     unetlayers(layers, outsz; skip_upscale = skip_upscale - 1,
+                                outplanes)...)
+    else
+        childunet = Chain(unetlayers(layers, outsz; skip_upscale)...)
+        outsz = Flux.outputsize(childunet, outsz)
+
+        inplanes = sz[end - 1]
+        midplanes = outsz[end - 1]
+        outplanes = isnothing(outplanes) ? inplanes : outplanes
+
+        return UNetBlock(Chain(layer, childunet),
+                         inplanes, midplanes, outplanes)
+    end
+end
+
+function UNetBlock(m_child, inplanes, midplanes, outplanes = 2 * inplanes)
+    return Chain(SkipConnection(Chain(m_child,
+                                      PixelShuffleICNR(midplanes, midplanes)),
+                                Parallel(cat_channels, identity, BatchNorm(inplanes))),
+                 xs -> relu.(xs),
+                 UNetCombineLayer(inplanes + midplanes, outplanes))
+end
+
+"""
+    unet(backbone; inputsize::NTuple{4, Integer}, outplanes::Integer = 3,
+    	final::Any = UNetFinalBlock, fdownscale::Integer = 0, kwargs...)
+
+Creates a UNet model with specified backbone. Backbone of Any Metalhead model
+can be used as encoder.
+([reference](https://arxiv.org/abs/1505.04597)).
+
+# Arguments
+
+    - `backbone`: The backbone layers to be used in encoder. 
+    	For example, `Metalhead.backbone(Metalhead.ResNet(18))` can be passed 
+    	to instantiate a UNet with layers of resnet18 as encoder.
+    - `inputsize`: size of input image
+    - `outplanes`: number of output feature planes
+    - `final`: final block as described in original paper
+    - `fdownscale`: downscale factor
+"""
+function unet(backbone, inputsize::NTuple{4, Integer}, outplanes::Integer,
+              final::Any = UNetFinalBlock, fdownscale::Integer = 0, kwargs...)
+    backbonelayers = collect(iterlayers(backbone))
+    layers = unetlayers(backbonelayers, inputsize; m_middle = UNetMiddleBlock,
+                        skip_upscale = fdownscale, kwargs...)
+
+    outsz = Flux.outputsize(layers, inputsize)
+    layers = Chain(layers, final(outsz[end - 1], outplanes))
+
+    return layers
+end
+
+"""
+    UNet(backbone, inputsize::NTuple{4, Integer}, outplanes::Integer = 3;
+        pretrain::Bool = false)
+
+Creates a UNet model with specified backbone. Backbone of Any Metalhead model can be used as
+encoder.
+([reference](https://arxiv.org/abs/1505.04597)).
+
+# Arguments
+
+    - `backbone`: The backbone layers to be used in encoder. 
+    	For example, `Metalhead.backbone(Metalhead.ResNet(18))` can be passed to instantiate a UNet with layers of
+    	resnet18 as encoder.
+    - `inputsize`: size of input image
+    - `outplanes`: number of output feature planes.
+    - `pretrain`: Whether to load the pre-trained weights for ImageNet
+
+!!! warning
+
+    `UNet` does not currently support pretrained weights.
+
+See also [`Metalhead.UNet`](@ref).
+"""
+struct UNet
+    layers::Any
+end
+@functor UNet
+
+function UNet(backbone, inputsize::NTuple{4, Integer}, outplanes::Integer = 3;
+              pretrain::Bool = false)
+    layers = unet(backbone, inputsize, outplanes)
+    if pretrain
+        loadpretrain!(layers, string("UNet"))
+    end
+    return UNet(layers)
+end
+
+(m::UNet)(x::AbstractArray) = m.layers(x)

src/utilities.jl

@@ -79,3 +79,7 @@ linear_scheduler(drop_prob::Nothing; depth::Integer) = fill(drop_prob, depth)
 function _checkconfig(config, configs)
     @assert config in configs "Invalid configuration. Must be one of $(sort(collect(configs)))."
 end
+
+# Utility function to return Iterator over layers, adopted from FastAI.jl
+iterlayers(m::Chain) = Iterators.flatten(iterlayers(l) for l in m.layers)
+iterlayers(m) = (m,)

test/convnets.jl

@@ -347,3 +347,11 @@ end
         _gc()
     end
 end
+
+@testset "UNet" begin
+    encoder = Metalhead.backbone(ResNet(18))
+    model = UNet(encoder, (256, 256, 3, 1), 10)
+    @test size(model(x_256)) == (256, 256, 10, 1)
+    @test gradtest(model, x_256)
+    _gc()
+end