lmflow.optim.lars

Classes

LARS	Extends SGD in PyTorch with LARS scaling from the paper

Module Contents

class lmflow.optim.lars.LARS(params, lr: float = 0.01, momentum: float = 0.0, dampening: float = 0.0, weight_decay: float = 0.0, nesterov: bool = False, trust_coefficient: float = 0.01, eps: float = 1e-08)

Bases: torch.optim.optimizer.Optimizer

Extends SGD in PyTorch with LARS scaling from the paper Large batch training of Convolutional Networks.

Note

The application of momentum in the SGD part is modified according to the PyTorch standards. LARS scaling fits into the equation in the following fashion.

\[\begin{split}\begin{aligned} g_{t+1} & = \text{lars_lr} * (\beta * p_{t} + g_{t+1}), \\ v_{t+1} & = \\mu * v_{t} + g_{t+1}, \\ p_{t+1} & = p_{t} - \text{lr} * v_{t+1}, \\end{aligned}\end{split}\]

where \(p\), \(g\), \(v\), \(\\mu\) and \(\beta\) denote the parameters, gradient, velocity, momentum, and weight decay respectively. The \(lars_lr\) is defined by Eq. 6 in the paper. The Nesterov version is analogously modified.

Warning

Parameters with weight decay set to 0 will automatically be excluded from layer-wise LR scaling. This is to ensure consistency with papers like SimCLR and BYOL.

Note

Reference code: PyTorchLightning/lightning-bolts

__setstate__(state) → None

step(closure=None)

Performs a single optimization step.

Parameters:

closure – A closure that reevaluates the model and returns the loss.