LRScheduler
On This Page
LRScheduler¶
class modelzoo.common.pytorch.optim.lr_scheduler.LRScheduler (optimizer, decay_steps: int = None, disable_lr_steps_reset: bool = False)
Cerebras specific learning rate scheduler base class.
The learning rate schedulers implemented in this file are specifically designed to be run on a Cerebras system. This means that there are certain caveats to these custom schedulers that differ from a typical LR scheduler found in core PyTorch.
The learning rate schedulers here are intended to be stepped at every iteration. This means lr_scheduler.step() should be called after every optimizer.step(). Hence, the learning rate schedulers operate on a step-by-step basis. Having said that, there are some variables used such as last_epoch that might indicate otherwise. The only reason these variables are used is to match what is used in core PyTorch. It does not indicate that things are operating on an epoch-by-epoch basis.
Also, note that the above means that our LR schedulers are incompatible with the LR schedulers found in core PyTorch. The state cannot simply be transferred between the two. So, one of the LR schedulers defined here must be used in order to have LR scheduling on the Cerebras system.
load_state_dict (state_dict:dict)¶
Loads the schedulers state.
When saving or loading the scheduler, make sure to also save or load the state of the optimizer.
Parameters:: state_dict (dict) scheduler state. Should be an object returned from a place to call to state_dict().
state_dict()¶
Returns the state of the scheduler as a dict.
It contains an entry for every variable in self.__dict__ which is not the optimizer. The learning rate lambda functions will only be saved if they are callable objects and not if they are functions or lambdas.
When saving or loading the scheduler, make sure to also save or load the state of the optimizer.