Revise `Optimizer` and add state-access methods.
Main idea :
The overall goal was to be able to checkpoint optimizers between rounds. In the process, we reviewed the existing Checkpointer class.
This MR has eventually been cut in two parts:
- the current MR retains only the
Optimizer-related revisions - another MR is being created for the
Checkpointer-related revisions
This MR now performs the following changes:
- Add a
get_stateandset_statepair of methods toOptiModuleandOptimizerthat enable accessing internal state variables. - Add the
Optimizer.start_roundmethod, that triggers wrappedRegularizer.start_roundones. Call it inTrainingManager.training_round. - Garbage-collect
OptiModule.(de)serializemethods, that have been made obsolete byOptiModule.from_specsand revisions to config methods. - Fix
YogiModule, that had previously been regressed into not deferring from its parentAdamModule. - Perform an overall clean-up of the
OptimizerandOptiModuleclasses and of their docstrings. - Implement unit tests for
Optimizer. - Update unit tests for
OptiModuleclasses.
Legacy to-do for the initial implementation by @nbigaud:
-
Create the abstract method for optimodule -
select the list of actual variables to be saved as part of this -
Implement it for the full list -
_adaptative.py:AdaGradModule.state,AdamModule.steps,AdamModule._vmax -
_momentum.py:MomentumModule.velocity,EWMAModule.state(and YogiMomentumModule inherits)
-
-
Create Optimizer-level methods to get and set state using json, which just stacks calls to the state method from modules -
Merge recent mainbranch changes
Edited by ANDREY Paul