brainpy.optim.CosineAnnealingLR

class brainpy.optim.CosineAnnealingLR(lr, T_max, eta_min=0.0, last_epoch=- 1)

Set the learning rate of each parameter group using a cosine annealing schedule, where \(\eta_{max}\) is set to the initial lr and \(T_{cur}\) is the number of epochs since the last restart in SGDR:

\[\begin{split}\begin{aligned} \eta_t & = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 + \cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right), & T_{cur} \neq (2k+1)T_{max}; \\ \eta_{t+1} & = \eta_{t} + \frac{1}{2}(\eta_{max} - \eta_{min}) \left(1 - \cos\left(\frac{1}{T_{max}}\pi\right)\right), & T_{cur} = (2k+1)T_{max}. \end{aligned}\end{split}\]

When last_epoch=-1, sets initial lr as lr. Notice that because the schedule is defined recursively, the learning rate can be simultaneously modified outside this scheduler by other operators. If the learning rate is set solely by this scheduler, the learning rate at each step becomes:

\[\eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 + \cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right)\]

It has been proposed in `SGDR: Stochastic Gradient Descent with Warm Restarts`_. Note that this only implements the cosine annealing part of SGDR, and not the restarts.

Parameters

• lr (float) – Initial learning rate.

• T_max (int) – Maximum number of iterations.

• eta_min (float) – Minimum learning rate. Default: 0.

• last_epoch (int) – The index of last epoch. Default: -1.

:param .. _SGDR: https://arxiv.org/abs/1608.03983 :type .. _SGDR: Stochastic Gradient Descent with Warm Restarts:

__init__(lr, T_max, eta_min=0.0, last_epoch=- 1)

Methods

__init__(lr, T_max[, eta_min, last_epoch])
cpu()	Move all variable into the CPU device.
cuda()	Move all variables into the GPU device.
load_state_dict(state_dict[, warn])	Copy parameters and buffers from state_dict into this module and its descendants.
load_states(filename[, verbose])	Load the model states.
nodes([method, level, include_self])	Collect all children nodes.
register_implicit_nodes(*nodes[, node_cls])
register_implicit_vars(*variables, ...)
save_states(filename[, variables])	Save the model states.
state_dict()	Returns a dictionary containing a whole state of the module.
step_call()
step_epoch()
to(device)	Moves all variables into the given device.
tpu()	Move all variables into the TPU device.
train_vars([method, level, include_self])	The shortcut for retrieving all trainable variables.
tree_flatten()	Flattens the object as a PyTree.
tree_unflatten(aux, dynamic_values)	New in version 2.3.1.
unique_name([name, type_])	Get the unique name for this object.
vars([method, level, include_self, ...])	Collect all variables in this node and the children nodes.

Attributes

name	Name of the model.