Flux.jlでSinGANを実装する

using Adapt
using BSON: @load, @save
using CuArrays
using Flux
using Flux: mse, pullback, glorot_normal
using Flux.Optimise: update!
using JSON
using OrderedCollections
using Random
using Statistics

expand_dim(dim...) = (dim..., 3, 1)
zeros_like(T::Type, dims...) = fill!(similar(T, dims...), 0f0)
zeros_like(xs::AbstractArray, dims...) = fill!(similar(xs, dims...), 0f0)
randn_like(T::Type, dims...) = randn!(similar(T, dims...))
randn_like(xs::AbstractArray, dims...) = randn!(similar(xs, dims...))

function size_pyramid(scale, min_size, image_size)
    current_size = min_size 
    pyramid = Vector{Tuple{Int64,Int64}}()
    for i in 1:100
        push!(pyramid, current_size)
        current_size == image_size && break
        current_size = @. floor(Int64, min_size * scale^i)
        current_size = min.(current_size, image_size)
    end
    return pyramid
end

channel_pyramid(n_stage) = min.(map(s->32 * 2^(floor(Int64, s / 4)), 1:n_stage), 128)

function resize_and_padding(x::Array{Float32,4}, 
            image_shape::Tuple{Int64,Int64}, padded_shape::Tuple{Int64,Int64})
    # println(size(x), image_shape, padded_shape)
    x_large = imresize(view(x, :, :, :, 1), image_shape...)
    xx = zeros(Float32, expand_dim(padded_shape...))
    pad1, pad2 = (@. (padded_shape - image_shape) ÷ 2)[1:2]
    xx[1 + pad1:image_shape[1] + pad1, 1 + pad2:image_shape[2] + pad2, : , 1] = x_large
    return xx
end

function resize_and_padding(x::CuArray{Float32,4}, 
            image_shape::Tuple{Int64,Int64}, padded_shape::Tuple{Int64,Int64})
    return cu(resize_and_padding(adapt(Array{Float32}, x), image_shape, padded_shape))
end

function build_image_pyramid(img::AbstractArray, image_shapes::Vector{Tuple{Int64,Int64}}, noise_shapes::Vector{Tuple{Int64,Int64}})
    return map((is, ns)->resize_and_padding(img, is, ns), image_shapes, noise_shapes)
end

function build_zero_pyramid(xs::AbstractArray, shapes::Vector{Tuple{Int64,Int64}})
    return map(s->zeros_like(xs, expand_dim(s...)), shapes)
end

function build_noise_pyramid(xs::AbstractArray, shapes::Vector{Tuple{Int64,Int64}}, amplifiers::Vector{Float32})
    return map((s, a)->a * randn_like(xs, expand_dim(s...)), shapes, amplifiers)
end

function build_rec_pyramid(xs::AbstractArray, shapes::Vector{Tuple{Int64,Int64}}, amplifier::Float32)
    v = build_zero_pyramid(xs, shapes)
    randn!(v[1])
    v[1] *= amplifier
    return v
end

# Re-define leakyrelu function
# https://github.com/FluxML/Flux.jl/issues/963
myleakyrelu(x::Real, a = oftype(x / one(x), 0.01)) = max(a * x, x / one(x))

conv_block(in, out) = [
        Conv((3, 3), in => out; init = glorot_normal, pad = (1, 1)),
        BatchNorm(out),
        x->myleakyrelu.(x, 0.2f0)
    ]

function build_layers(n_layers, in_chs, conv_chs, out_chs, σ)
    layers = conv_block(in_chs, conv_chs)
    for _ in 1:n_layers - 2
        push!(layers, conv_block(conv_chs, conv_chs)...)
    end
    tail_layer = Conv((3, 3), conv_chs => out_chs, σ;
        init = glorot_normal, pad = (1, 1))
    push!(layers, tail_layer)
    return Chain(layers...)
end

"""
    DiscriminatorPyramid
"""
mutable struct DiscriminatorPyramid{T <: Tuple}
    chains::T
    DiscriminatorPyramid(xs...) = new{typeof(xs)}(xs)
end

build_single_discriminator(n_layers, conv_chs) = build_layers(n_layers, 3, conv_chs, 1, identity)

function DiscriminatorPyramid(n_stage::Integer, n_layers::Integer)
    ds = build_single_discriminator.(n_layers, channel_pyramid(n_stage))
    return DiscriminatorPyramid(gpu.(ds)...)
end

function DiscriminatorPyramid(image_shapes::Vector{Tuple{Int64,Int64}}, n_layers::Integer)
    DiscriminatorPyramid(Base.length(image_shapes), n_layers)
end

function Base.show(io::IO, d::DiscriminatorPyramid)
    print(io, "DiscriminatorPyramid(")
    join(io, d.chains, ", \n")
    print(io, ")")
end

"""
    NoiseConnection
"""
mutable struct NoiseConnection
    layers
    pad::Int64
end

@Flux.functor NoiseConnection

function (nc::NoiseConnection)(prev::T, noise::T) where {T <: AbstractArray{Float32,4}}
    pad = nc.pad
    raw_output = nc.layers(noise + prev)::T + prev
    return raw_output[1 + pad:end - pad, 1 + pad:end - pad, :, :]
end

function Base.show(io::IO, nc::NoiseConnection)
    print(io, "NoiseConnection(", nc.layers, ", ", nc.pad, ")")
end

"""
    GeneratorPyramid
"""
mutable struct GeneratorPyramid{T <: Tuple}
    image_shapes::Vector{Tuple{Int64,Int64}}
    noise_shapes::Vector{Tuple{Int64,Int64}}
    pad::Int64
    chains::T
    GeneratorPyramid(image_shapes, noise_shapes, pad, xs...) = new{typeof(xs)}(image_shapes, noise_shapes, pad, xs)
end

build_single_gen_layers(n_layers, conv_chs) = build_layers(n_layers, 3, conv_chs, 3, tanh)
build_single_generator(n_layers, conv_chs, pad) = NoiseConnection(build_single_gen_layers(n_layers, conv_chs), pad)

function GeneratorPyramid(image_shapes::Vector{Tuple{Int64,Int64}}, n_layers::Integer, pad::Integer = 5)
    n_stage = Base.length(image_shapes)
    # receptive field = 11, floor(11/2) = 5
    noise_shapes = [2 * pad .+ s for s in image_shapes]
    ds = build_single_generator.(n_layers, channel_pyramid(n_stage), pad)
    return GeneratorPyramid(image_shapes, noise_shapes, pad, gpu.(ds)...)
end

function Base.show(io::IO, d::GeneratorPyramid)
    print(io, "GeneratorPyramid(")
    print(io, d.image_shapes, ", ")
    print(io, d.noise_shapes, ", ")
    println(io, d.pad, ", ")
    join(io, d.chains, ", \n")
    print(io, ")")
end

function (genp::GeneratorPyramid)(xs::AbstractVector{T}, st::Integer, resize::Bool) where {T <: AbstractArray{Float32,4}}
    if st == 0
        zeros_shape = resize ? first(genp.noise_shapes) : first(genp.image_shapes)
        return zeros_like(T, expand_dim(zeros_shape...))
    end
    prev = genp(xs, st - 1, true)
    out = genp.chains[st](prev, xs[st])
    return resize ? resize_and_padding(out, genp.image_shapes[st + 1], genp.noise_shapes[st + 1]) : out
end

function discriminator_loss(d_real::AbstractArray, d_g_fake_adv::AbstractArray)
    real_loss = mse(1f0, mean(d_real; dims = (1, 2)))
    fake_loss = mse(0f0, mean(d_g_fake_adv; dims = (1, 2)))
    return real_loss + fake_loss
end

generator_adv_loss(d_g_fake_adv::AbstractArray) = mse(1f0, mean(d_g_fake_adv; dims = (1, 2)))

generator_rec_loss(real_img::AbstractArray, g_fake_rec::AbstractArray) = mse(real_img, g_fake_rec)

function update_discriminator!(opt, dscr, real_img, g_fake_adv)
    @eval Flux.istraining() = true
    ps = params(dscr)
    grad = gradient(ps) do
        discriminator_loss(dscr(real_img), dscr(g_fake_adv))
    end
    update!(opt, ps, grad)
    @eval Flux.istraining() = false
end

function update_generator!(opt, dscr, gen, real_img, prev_rec, prev_adv, noise_rec, noise_adv, alpha)
    @eval Flux.istraining() = true
    ps = params(gen)
    grad = gradient(ps) do
        g_fake_rec = gen(prev_rec, noise_rec)
        d_g_fake_adv = dscr(gen(prev_adv, noise_adv))
        generator_adv_loss(d_g_fake_adv) + alpha * generator_rec_loss(real_img, g_fake_rec)
    end
    update!(opt, ps, grad)
    @eval Flux.istraining() = false
end

grad = gradient(ps) do
    discriminator_loss(dscr(real_img), dscr(g_fake_adv))
end

grad = gradient(() -> discriminator_loss(dscr(real_img), dscr(g_fake_adv)), ps)

function update_discriminator!(opt, dscr, real_img, g_fake_adv)
    @eval Flux.istraining() = true
    ps = params(dscr)
    loss, back = pullback(ps) do
        discriminator_loss(dscr(real_img), dscr(g_fake_adv))
    end
    grad = back(Zygote.sensitivity(loss))
    update!(opt, ps, grad)
    @eval Flux.istraining() = false
    return loss
end

function train_epoch!(opt_dscr, opt_gen, st, loop_dscr, loop_gen,
        dscr, genp, prev_rec, noise_rec, real_img, amplifiers, alpha)

    # discriminator
    foreach(1:loop_dscr) do _
        noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
        g_fake_adv = genp(noise_adv, st, false)
        update_discriminator!(opt_dscr, dscr, real_img, g_fake_adv)
    end

    # generator
    foreach(1:loop_gen) do _
        noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
        prev_adv = genp(noise_adv, st - 1, true)
        update_generator!(opt_gen, dscr, genp.chains[st], real_img, prev_rec, prev_adv, noise_rec, last(noise_adv), alpha)
    end

    noise_adv = build_noise_pyramid(prev_rec, genp.noise_shapes[1:st], amplifiers)
    g_fake_adv = genp(noise_adv, st, false)
    loss_dscr = discriminator_loss(dscr(real_img), dscr(g_fake_adv))
    d_g_fake_adv = dscr(g_fake_adv)
    loss_gen_adv = generator_adv_loss(d_g_fake_adv)
    g_fake_rec = genp.chains[st](prev_rec, noise_rec)
    loss_gen_rec = generator_rec_loss(real_img, g_fake_rec) 

    return loss_dscr, loss_gen_adv, loss_gen_rec
end

mutable struct HyperParams
    scale::Float64                  # progression scale, > 1
    min_size_x::Int64               # minimal image width
    min_size_y::Int64               # minimal image height
    img_size_x::Int64               # output image width
    img_size_y::Int64               # output image height
    n_layers::Int64                 # number of conv layers
    max_epoch::Int64                # training epochs
    reduce_lr_epoch::Int64          # reduce learining rate after training `redule_lr_epoch` epochs
    save_image_every_epoch::Int64   # save generated image every `save_image_every_epoch` epoch
    save_loss_every_epoch::Int64    # save loss every `save_loss_every_epoch` epoch
    loop_dscr::Int64                # training steps par descriminator training epoch
    loop_gen::Int64                 # training steps par generator training epoch
    lr_dscr::Float64                # discriminator learining rate
    lr_gen::Float64                 # generator learning rate
    alpha::Float32                  # rec loss coefficient
    amplifier_init::Float32         # noise amplifier
    HyperParams() = new(4/3, 25, 25, 128, 128, 5, 2000, 1600, 500, 100, 3, 3, 5e-4, 5e-4, 50f0, 1f0)
end

show_dict(hp::HyperParams) = OrderedDict(string(nm) => getfield(hp, nm) for nm in fieldnames(HyperParams))
image_shapes(hp::HyperParams) = size_pyramid(hp.scale, (hp.min_size_x, hp.min_size_y), (hp.img_size_x, hp.img_size_y))

function setup_models(hp::HyperParams)
    img_shapes = image_shapes(hp)
    dscrp = DiscriminatorPyramid(img_shapes, hp.n_layers) |> gpu
    genp = GeneratorPyramid(img_shapes, hp.n_layers) |> gpu
    return dscrp, genp
end

function estimate_noise_amplifier(prev_rec::AbstractArray{Float32,4}, real_img::AbstractArray{Float32,4},
        pad::Integer, amplifier_init::Float32)
    prev_rec_crop = @view prev_rec[1 + pad:end - pad, 1 + pad:end - pad, :, :]
    rmse = sqrt(mse(real_img, prev_rec_crop))
    return rmse * amplifier_init
end

function train!(dscrp::DiscriminatorPyramid, genp::GeneratorPyramid, 
        real_img_p::Vector{T}, hp::HyperParams) where {T <: AbstractArray{Float32,4}}
    stages = Base.length(genp.image_shapes)

    amplifiers = Float32[]

    # fixed noise for rec
    fixed_noise_rec = build_rec_pyramid(first(real_img_p), genp.noise_shapes, 1f0)
    fixed_noise_adv = similar(fixed_noise_rec)
    
    for st in 1:stages
        @info "Step $(st)"
        # reset optimizer
        opt_dscr = ADAM(hp.lr_dscr, (0.5, 0.999))
        opt_gen = ADAM(hp.lr_gen, (0.5, 0.999))

        # calculate noise amplifier
        prev_rec = genp(fixed_noise_rec, st - 1, true) # padded
        amp = st == 1 ? 1f0 : estimate_noise_amplifier(prev_rec, real_img_p[st], genp.pad, hp.amplifier_init)
        push!(amplifiers, amp)
        # add noise for adv 
        fixed_noise_adv[st] = amp * randn_like(prev_rec, expand_dim(genp.noise_shapes[st]...))

        save_noise_amplifiers(st, amp)
        @info "Noise amplifier = $(amp)"

        @time for ep in 1:hp.max_epoch
            # reduce learnint rate
            if ep == hp.reduce_lr_epoch
                @info "Reduce learning rate"
                opt_dscr.eta /= 10
                opt_gen.eta /= 10
            end

            loss_dscr, loss_gen_adv, loss_gen_rec =
                train_epoch!(opt_dscr, opt_gen, st, hp.loop_dscr, hp.loop_gen,
                    dscrp.chains[st], genp, prev_rec, fixed_noise_rec[st], real_img_p[st], amplifiers, hp.alpha)
        end
    end

end

$ convert -delay 10 -loop 0 img*.png fluxjl-singan_lightning.gif