Open Vilin97 opened 11 months ago
One possible design is this:
reset!(tree) = foreach(reset!, tree)
reset!(ℓ::Leaf) = ℓ.state = reset!(ℓ.rule, ℓ.state)
reset!(::AbstractRule, ::Nothing) = nothing
reset!(rule::AbstractRule, state) = throw(ArgumentError("""reset! does not now how to handle this rule."))
Then rules need to opt-in by defining a method of 2-arg reset!
... with some fill!!
which allows for immutable arrays?
reset!(rule::Adam, (mt, vt, βt)) = (fill!!(mt, 0), fill!!(vt, 0), rule.beta)
We can't easily fall back to calling init
again for unknown rules, as we don't have the original parameters x
here.
Falling back to zero like this might be OK for built-in rules like Momentum etc, but could be wrong for user-defined rules... probably we shouldn't:
reset!(rule::AbstractRule, state::AbstractArray) = fill!!(state, 0)
We could always make reset!
take the parameter tree of x
s too, but that may come at the cost of sacrificing type stability.
Motivation and description
In my application I do 25 steps of gradient descent
update!
steps in a loop (solving a differential equation). I need the momentum from the previous 25 GD steps to NOT carry over to the next 25 GD steps. In other words, the behavior I am looking for is analogous to callingFlux.setup(optimiser, model)
every time. Unfortunately,Flux.setup
is type-unstable https://github.com/FluxML/Optimisers.jl/issues/162. It would be great to have a functionreset!(optimiser_state)
that resets the momenta. Maybe a more stringent requirement is thatholds.
Possible Implementation
Below is an implementation for
Adam
.