Add support for Optimisers.jl #114
Conversation
src/training.jl
Outdated
| @@ -49,18 +49,29 @@ function step! end | |||
| function step!(learner, phase::TrainingPhase, batch) | |||
| xs, ys = batch | |||
| runstep(learner, phase, (; xs=xs, ys=ys)) do handle, state | |||
| state.grads = gradient(learner.params) do | |||
| state.ŷs = learner.model(state.xs) | |||
| state.grads, _, _= gradient(learner.model, state.xs, state.ys) do model, xs, ys | |||
There was a problem hiding this comment.
I'm not sure if passing in all three is necessary here? Since we want the gradients of the model with respect to xs and ys, but I wonder if this calculates some unneeded gradients the other way around as well?
src/training.jl
Outdated
| end | ||
| end | ||
|
|
||
| # Handle both old Flux.jl and new Optimisers.jl optimisers | ||
| function _update!(optimizer::Flux.Optimise.AbstractOptimiser, params, model, grads) | ||
| update!(optimizer, model, grads) |
There was a problem hiding this comment.
This currently throws an error. For context params isa Params and grads is no longer a Grads. Is a Params even needed anymore?
There was a problem hiding this comment.
Left a comment above about this.
|
|
||
|
|
||
| @testset "Optimisers.jl compatibility" begin | ||
| learner = testlearner(coeff = 3, opt=Optimisers.Descent(0.001)) |
There was a problem hiding this comment.
This already passes 👍 but all the old optim tests are broken for the time being
src/training.jl
Outdated
| state.grads = gradient(learner.params) do | ||
| state.ŷs = learner.model(state.xs) | ||
|
|
||
| state.grads, = gradient(learner.model) do model |
There was a problem hiding this comment.
I think here you want to take the gradient w.r.t. learner.params when the optimizer is a Flux.Optimise.AbstractOptimiser. Conversely, if it is not, you take the gradient w.r.t. learner.model like you are now.
This is why update! below is erroring cause you need to Grads object for the old optimizers. And you can only get that with implicit params.
I think some dispatch for the gradient would be easiest. Another option is to have a utility that takes the model, the gradient w.r.t. it, and Params, then it produces a Grads to match.
There was a problem hiding this comment.
Ah, I figured. Was hoping there may be a way to have the same Zygote.gradient call work but I guess not. I'll add a dispatch on the optimiser there.
There was a problem hiding this comment.
Well what you're asking for might come eventually in a later version of Flux as part of the AD-agnostic push. So, the code might eventually get simpler.
There was a problem hiding this comment.
That's good to know. Definitely let me know then, so I can clean this up again.
src/training.jl
Outdated
| end | ||
| end | ||
|
|
||
| # Handle both old Flux.jl and new Optimisers.jl optimisers | ||
| function _update!(optimizer::Flux.Optimise.AbstractOptimiser, params, model, grads) | ||
| update!(optimizer, model, grads) |
There was a problem hiding this comment.
Left a comment above about this.
|
@darsnack I added the dispatch for |
Closes #112 (once done). @ToucheSir @darsnack
So this is a first draft for adding Optimisers.jl support (new optims) while keeping compatibility with optimisers in
Flux.Optimise(old optims).Passing in new optims already works, but I've broken support for old optims. Before, FluxTraining.jl was using implicit parameters with
ParamsandGradsobjects. I'm not sure how to use the old optims with explicit parameters togradient.I'll leave some more questions next to the code changes, some feedback from you two would be much appreciated!