
using Flux


using CUDA


using ParameterSchedulers


using MLDatasets


using Images


using Augmentor


using IterTools


using OnlineStats


using Printf


Base.@kwdef struct Config
    batchsize = 32
    throttle = 40
    lr = 1e-3
    wd = 1f-2
    epochs = 100
end

config = Config();


if CUDA.has_cuda()
    @info "CUDA is on"
    CUDA.allowscalar(false)
end

┌ Info: CUDA is on
└ @ Main In[11]:2


# CIFAR10.download(i_accept_the_terms_of_use=true)


train_data = CIFAR10.traindata(Float32)
test_data = CIFAR10.testdata(Float32);


train_aug = FlipX(0.5) |> ShearX(-5:5) * ShearY(-5:5) |> Rotate(-15:15) |> CropSize(32,32) |>
            Zoom(0.9:0.1:1.2) |> SplitChannels() |> PermuteDims(3, 2, 1) |> ConvertEltype(Float32)

8-step Augmentor.ImmutablePipeline:
 1.) Either: (50%) Flip the X axis. (50%) No operation.
 2.) Either: (50%) ShearX by ϕ ∈ -5:5 degree. (50%) ShearY by ψ ∈ -5:5 degree.
 3.) Rotate by θ ∈ -15:15 degree
 4.) Crop a 32×32 window around the center
 5.) Zoom by I ∈ {0.9×0.9, 1.0×1.0, 1.1×1.1, 1.2×1.2}
 6.) Split colorant into its color channels
 7.) Permute dimension order to (3, 2, 1)
 8.) Convert eltype to Float32


function collate((imgs, labels))
    imgs = imgs |> gpu
    labels = Flux.onehotbatch(labels .+ 1, 1:10) |> gpu
    imgs, labels
end

collate (generic function with 1 method)


function collate((imgs, labels), aug)
    imgs_aug = Array{Float32}(undef, size(imgs))
    augmentbatch!(imgs_aug, CIFAR10.convert2image(imgs), aug)
    collate((imgs_aug, labels))
end

collate (generic function with 2 methods)


train_loader = imap(d -> collate(d, train_aug),
    Flux.Data.DataLoader(train_data, batchsize=config.batchsize, shuffle=true))

test_loader = imap(collate,
    Flux.Data.DataLoader(test_data, batchsize=config.batchsize, shuffle=false));


batch = iterate(train_loader)[1]
println(CIFAR10.classnames()[Flux.onecold(cpu(batch[2])[:, 1], 1:10)])
CIFAR10.convert2image(cpu(batch[1])[:,:,:, 1])

truck


function conv_block(ch::Pair; kernel_size=3, stride=1, activation=relu)
    Chain(Conv((kernel_size, kernel_size), ch, pad=SamePad(), stride=stride, bias=false, init=Flux.kaiming_normal),
          BatchNorm(ch.second, activation))
end

conv_block (generic function with 1 method)


function basic_residual(ch::Pair)
    Chain(conv_block(ch),
          conv_block(ch.second => ch.second, activation=identity))
end

basic_residual (generic function with 1 method)


struct AddMerge
    gamma
    expand
end

Flux.@functor AddMerge

function AddMerge(ch::Pair)
    if ch.first == ch.second
        expand = identity
    else
        expand = conv_block(ch, kernel_size=1, activation=identity)
    end
    AddMerge([0.f0], expand)
end

(m::AddMerge)(x1, x2) = relu.(m.gamma .* x1 .+ m.expand(x2))


function residual_block(ch::Pair)
    residual = basic_residual(ch)
    SkipConnection(residual, AddMerge(ch))
end

residual_block (generic function with 1 method)


function residual_stack(in_channels, repetitions, downsamplings)
    layers = []
    res_channels = in_channels
    for (rep, stride) in zip(repetitions, downsamplings)
        if stride > 1
            push!(layers, MaxPool((stride, stride)))
        end
        for i = 1:rep
            push!(layers, residual_block(in_channels => res_channels))
            in_channels = res_channels
        end
        res_channels *= 2
    end
    Chain(layers...)
end

residual_stack (generic function with 1 method)


function stem(in_channels=3; channel_list = [32, 32, 64], stride=1)
    layers = []
    for channels in channel_list
        push!(layers, conv_block(in_channels => channels, stride=stride))
        in_channels = channels
        stride=1
    end
    Chain(layers...)
end

stem (generic function with 2 methods)


function head(in_channels, classes, p_drop=0.)
    Chain(GlobalMeanPool(),
          flatten,
          Dropout(p_drop),
          Dense(in_channels, classes))
end

head (generic function with 2 methods)


function resnet(classes, repetitions, downsamplings; in_channels=3, p_drop=0.)
    Chain(stem(in_channels, stride=downsamplings[1]),
          residual_stack(64, repetitions, downsamplings[1:end]),
          head(64 * 2^(length(repetitions) - 1), classes, p_drop))
end

resnet (generic function with 1 method)


model = resnet(10, [2, 2, 2, 2], [1, 1, 2, 2, 2], p_drop=0.3) |> gpu;


mutable struct AdaptiveWeightDecay{O}
    opt::O
    wd::Real
end

function Flux.Optimise.apply!(o::AdaptiveWeightDecay, x, Δ)
    opt = o.opt
    Δ = Flux.Optimise.apply!(opt, x, Δ)
    @. Δ += o.wd * opt.eta * x
    return Δ
end


FixedADAMW(η = 0.001, β = (0.9, 0.999), decay = 0) = AdaptiveWeightDecay(Flux.ADAM(η, β), decay)

FixedADAMW (generic function with 4 methods)


struct Cos2{T, S<:Integer}
    r_start::T
    r_end::T
    period::S
end

(schedule::Cos2)(t) = schedule.r_end + (schedule.r_start - schedule.r_end) * (1 + cos(π * (t - 1) / schedule.period)) / 2

Base.iterate(schedule::Cos2, t = 1) = schedule(t), t + 1


struct SchedPair
    sched1
    sched2
end

(schedule::SchedPair)(t) = (schedule.sched1(t), schedule.sched2(t))

Base.iterate(schedule::SchedPair, t = 1) = schedule(t), t + 1


function one_cycle(max_lr, total_steps; pct_start=0.3, div_factor=25., final_div_factor=1e4, max_m=0.95, base_m=0.85)
    initial_lr = max_lr / div_factor
    min_lr = initial_lr / final_div_factor
    steps_up = floor(Int, pct_start * total_steps)
    steps_down = total_steps - steps_up
    SchedPair(
        Sequence([Cos2(initial_lr, max_lr, steps_up), Cos2(max_lr, min_lr, steps_down)], [steps_up, steps_down]),
        Sequence([Cos2(max_m, base_m, steps_up), Cos2(base_m, max_m, steps_down)], [steps_up, steps_down]))
end

one_cycle (generic function with 1 method)


function update_func(o::AdaptiveWeightDecay{Flux.Optimise.ADAM}, s::Tuple{Float64,Float64})
    opt = o.opt
    opt.eta = s[1]
    opt.beta = (s[2], opt.beta[2]) # Tuple is immutable, therefore it needs to be replaced entirely
end

update_func (generic function with 1 method)


sched = one_cycle(config.lr, config.epochs * length(train_loader));


opt = ParameterSchedulers.Scheduler(sched, FixedADAMW(config.lr, (0.9, 0.999), config.wd), update_func);


loss(x, y) = Flux.logitcrossentropy(model(x), y);


ps = params(model);


function accuracy(model, data)
    m = Mean()
    for (x, y) in data
        fit!(m, Flux.onecold(cpu(model(x)), 1:10) .== Flux.onecold(cpu(y), 1:10))
    end
    value(m)
end

accuracy (generic function with 1 method)


evalcb = Flux.throttle(config.throttle) do
    @printf "Val accuracy: %.3f\n" accuracy(model, test_loader)
end;


Flux.@epochs config.epochs Flux.train!(loss, ps, train_loader, opt, cb=evalcb)

┌ Info: Epoch 1
└ @ Main /home/julius/.julia/packages/Flux/qp1gc/src/optimise/train.jl:135

Val accuracy: 0.100
Val accuracy: 0.426

┌ Info: Epoch 2
└ @ Main /home/julius/.julia/packages/Flux/qp1gc/src/optimise/train.jl:135

Val accuracy: 0.501
Val accuracy: 0.538

┌ Info: Epoch 3
└ @ Main /home/julius/.julia/packages/Flux/qp1gc/src/optimise/train.jl:135

Val accuracy: 0.570
Val accuracy: 0.573

┌ Info: Epoch 4
└ @ Main /home/julius/.julia/packages/Flux/qp1gc/src/optimise/train.jl:135

Val accuracy: 0.625
Val accuracy: 0.659

┌ Info: Epoch 5
└ @ Main /home/julius/.julia/packages/Flux/qp1gc/src/optimise/train.jl:135

Val accuracy: 0.675
Val accuracy: 0.698

ResNet on CIFAR10 using Julia¶

Configuration¶

Data¶

Model¶

Training¶

AdamW as in PyTorch¶

1Cycle scheduler¶

Start Training¶