add step! #1833
add step! #1833
Conversation
|
Similar to #1017 which should be complete |
|
True, sorry I missed that. Let's just move ahead with this one since it has a docstring and Felix put in some time recently. |
There was a problem hiding this comment.
A couple of docstring changes. The implementation looks good to me.
Can you also update the docs for "Custom Training loops" to use this new function? And add the docstring into the docs in that location?
One other suggestion is to call this trainstep!, since step! is very generic, and we might want to hold onto the name for something else.
Co-authored-by: Kyle Daruwalla <daruwalla.k.public@icloud.com>
Co-authored-by: Kyle Daruwalla <daruwalla.k.public@icloud.com>
| @@ -81,29 +81,29 @@ batchmemaybe(x) = tuple(x) | |||
| batchmemaybe(x::Tuple) = x | |||
|
|
|||
| """ | |||
| step!(loss, params, opt) | |||
| optimstep!(loss, params, opt) | |||
There was a problem hiding this comment.
I suggest optimstep! instead of trainstep! to indicate that this is the optimiser interface and keep the ML jargon to a minimum
There was a problem hiding this comment.
One vote for something evoking train! to stress that they are closely related.
If the longer-term plan is to use Optimisers.jl, this may not fit with train! at all -- some recent discussion here: #1902 (comment) . In which case there will be an implicit-style train! & Params story, and an explicit-style gradient and Optimisers.update!. With such a divide, this function wants to be clearly on the train! & Params side.
Maybe it should just be 3-arg train!? Without a data iterator, there is no iteration, that's all:
train!(loss, ::Params, data, ::AbstractOptimiser) # calls loss(d...) for d in data
train!(loss, ::Params, ::AbstractOptimiser) # calls loss() since there is no data
| # Calculate the gradients of the parameters | ||
| # with respect to the loss function | ||
| grads = Flux.gradient(parameters) do | ||
| # Update the parameters based on the chosen |
There was a problem hiding this comment.
This is right at the beginning instead of in the Custom Training Loop Section. It seems to me like the custom training loop section might either be redundant or demonstrate how to have a custom gradient calculation now.
Co-authored-by: Brian Chen <ToucheSir@users.noreply.github.com>
Co-authored-by: Brian Chen <ToucheSir@users.noreply.github.com>
Add
step!like suggested in #666 as a single step oftrain!to allow for more exotic optimisers to simply overloadstep!and be able to use thetrain!wrapper.PR Checklist