ande.jl

using Knet
using CUDA

setoptim!(m, optimizer) = for p in params(m); p.opt = Knet.clone(optimizer); end

dice(x, y; smooth::Float32=1.f0) = (2*sum(y .* x) + smooth) / (sum(y.^2) + sum(x.^2) + smooth)
loss(x, y) = 1 - dice(x, y)

function minimize!(model, x::KnetArray, y::KnetArray)
    ld = @diff loss(model(x), y)
    for w in params(model)
        Knet.update!(w, grad(ld, w))
    end
    return value(ld)
end

# Define a chain of layers :
struct Chain; layers; end
(c::Chain)(x) = (for l in c.layers; x = l(x); end; x)

struct test_model; c; end
function (m::test_model)(x)
    x = m.c(x)
    return x
end

function test_model()
    w = param(3, 3, 3, 1, 8)
    c = Chain((
               x->conv4(w, x, stride=2, padding=1),
               x->unpool(x),
               x->conv4(w, x, stride=2, padding=1),
               x->unpool(x),
               x->conv4(w, x, stride=2, padding=1),
               x->unpool(x)
              ))
    test_model(c)
end

# Main training loop
function main(N=500,T=1)

    # Get model
    model = test_model()
    setoptim!(model, Adam())

    # Kick off the training loop
    for i in 1:T
        @info "Epoch $i of $T"

        for i in 1:5
            x = rand(Float32, N, 256, 256, 1, 1)
            y = rand(Float32, N, 256, 256, 1, 1)
            train_loss = minimize!(model, KnetArray(x), KnetArray(y))
            CUDA.memory_status()
        end

        println("")
    end
end

main()