Flux unet

unet.jl

using Flux

downsample(c_in, c_out) = Chain(
    Conv((4, 4), c_in => c_out, stride = 2, pad = 1),
    BatchNorm(c_out, relu)
)

upsample(c_in, c_out) = Chain(
    ConvTranspose((4, 4), c_in => c_out, stride = 2, pad = 1),
    BatchNorm(c_out, relu)
)

function UNet(initial_channels, depth)
    down_path = map(1:depth) do d
        c_in = initial_channels *  2 ^ (d - 1)
        downsample(c_in, 2 * c_in)
    end

    up_path = map(depth:-1:1) do d
        c_in = initial_channels * 2 ^ (d + 1)
        upsample(d == depth ? c_in ÷ 2 : c_in, c_in ÷ 4)
    end

    input -> begin
        intermediate = TrackedArray[]

        foreach(down_path) do layer
            input = layer(input)
            push!(intermediate, input)
        end

        output = first(up_path)(pop!(intermediate))

        foreach(up_path[2:end]) do layer
            output = layer(cat(output, pop!(intermediate), dims = 3))
        end

        output
    end
end

initial_channels = 16
depth = 5

model = UNet(initial_channels, depth)
pars = params((model.down_path..., model.up_path...))

input = rand(Float32, 128, 128, initial_channels, 1)
output = model(input)

gs = gradient(pars) do
    sum(abs, input - model(input))
end