Example of how to use with Optimisers.jl

[theabhirath]

Hi, I've been trying to use this package with Optimisers.jl (specifically, I've been trying to use a Step schedule with a Scheduler but I seem to be getting errors that suggest that this setup works with the Flux optimisers, and not with Optimisers.jl for now. Is there a way to write code that works with Optimisers.jl?

[darsnack]

Constructing optimizers from Optimisers.jl is cheap and simple, since the state is de-coupled. Something like this would work:

struct Scheduler{T, F}
    constructor::F
    schedule::T
end

_get_opt(scheduler::Scheduler, t) = scheduler.constructor(scheduler.schedule(t))

Optimisers.init(o::Scheduler, x::AbstractArray) =
    (t = 1, opt = Optimisers.init(_get_opt(o, 1), x))

function Optimisers.apply!(o::Scheduler, state, x, dx)
    opt = _get_opt(o, state.t)
    new_state, new_dx = Optimisers.apply!(opt, state.opt, x, dx)

    return (t = state.t + 1, opt = new_state), new_dx
end

opt = Scheduler(Step(init_lr, decay)) do lr
    Momentum(lr)
end
st = Optimisers.setup(opt, model)

[theabhirath]

I tried this and my model stopped training 😬 It's stuck after one epoch

[darsnack]

That's weird. Which rule are you using? And can you post the value of st after setup for a small model? I would also print opt inside the definition for apply!.

[darsnack]

I would also confirm that training doesn't stall with a reasonable fixed LR first.

[theabhirath]

Whoops, nevermind, figured it out. I set step_sizes = 25 for Step, but naturally since Step is being called every step and not every epoch step_sizes has to be (25 * dataset_size) / batch_size. Reducing it by 0.1 every step meant that the learning rate was on the order of 10^-40 before even starting the second epoch 😅 It's training now, thank you!

[darsnack]

Ah good to know. You should check out Interpolator and the docs for it (under complex schedules). Just a slightly less error prone/cleaner way of specifying schedules in epochs that will be iterated by mini-batch.

[ToucheSir]

Now that we have Optimisers.adjust!, should Scheduler be modernized and adopted as the recommended (easy) way to schedule parameters with Flux + Optimisers.jl?

[Qiyu-Zh]

I meet a problems that scheduler can not be setup. Could you share me more details?

[darsnack]

The details for calling setup are in the original code above. Can you share the error?

[Qiyu-Zh]

[darsnack]

Modify the original code as

struct Scheduler{T, F} <: Optimisers.AbstractRule
    constructor::F
    schedule::T
end

[Qiyu-Zh]

Great, thanks!

[Qiyu-Zh]

Why there is the question the EXP not compatible with float?