Learning rate decay in callback function

[KianH]

I was wondering if there is a way to access the learning rate through a callback function when using Lux.jl and Optimization.jl packages. For instance in the line below:

Optimization.solve(optprob, ADAM(1e-3), callback = callback, progress = true, maxiters = sw)

I have set the learning rate to 1e-3, and I would like to access it through the callback function to use learning rate decay. I have found the following relevant resources:

https://fluxml.ai/Optimisers.jl/dev/#Adjusting-Hyperparameters

How to update learning rate during Flux training in a better manner?

Hi, I am trying to update the training rate during training. I did it with custom training loop like below: using Flux using Flux: @epochs using Flux: Flux.Data.DataLoader M = 10 N = 15 O = 2 X = repeat(1.0:10.0, outer=(1, N)) #input Y = repeat(1.0:2.0, outer=(1, N)) #output data = DataLoader(X,Y, batchsize=5, shuffle=true) dims = [M, O] layers = [Dense(dims[i], dims[i+1]) for i in 1:length(dims)-1] m = Chain(layers...) L(x, y) = Flux.Losses.mse(m(x), y) #cost function ps = Flux.par…

However both resources are intended for training their model using a loop, while the training for my code happens in one line, which is the line of code I pasted above. Therefore I would need to access the learning rate through a callback function. Do you have any suggestions?

[Vaibhavdixit02]

It should be possible but needs a change in the OptimizationOptimisers wrapper, we currently don’t pass the state as an argument to the callback function but it would be a small change and will then let you do it how it is described in the Flux docs. Can you open an issue in Optimization.jl and I’ll create a PR

[KianH]

In the documentations of Flux.jl regarding learning rate decay, they create a Flux.setup which is then used to update throughout the training as described in the link below.

https://fluxml.ai/Flux.jl/stable/training/training/

The optimization part of my code is written as the following:

adtype = Optimization.AutoZygote()

optf = Optimization.OptimizationFunction((x, p) -> loss(x,u0), adtype)

optprob = Optimization.OptimizationProblem(optf, ComponentArray(p))

opt = Adam(1e-2)

res1 = Optimization.solve(optprob, opt, callback = callback)

Meaning that I do not have any Flux.setup. Additionally, I am using Lux.jl.

What is equivalent to opt_state in the code snippet I have provided above to be able to perform learning rate decay through Lux.adjust!(…)?

[Vaibhavdixit02]

This works now Optimization.jl/lib/OptimizationOptimisers/test/runtests.jl at master · SciML/Optimization.jl · GitHub