File size: 3,098 Bytes
a256709 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
"""NovoGrad Optimizer.
Original impl by Masashi Kimura (Convergence Lab): https://github.com/convergence-lab/novograd
Paper: `Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks`
- https://arxiv.org/abs/1905.11286
"""
import torch
from torch.optim.optimizer import Optimizer
import math
class NovoGrad(Optimizer):
def __init__(
self,
params,
grad_averaging=False,
lr=0.1,
betas=(0.95, 0.98),
eps=1e-8,
weight_decay=0,
):
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
super(NovoGrad, self).__init__(params, defaults)
self._lr = lr
self._beta1 = betas[0]
self._beta2 = betas[1]
self._eps = eps
self._wd = weight_decay
self._grad_averaging = grad_averaging
self._momentum_initialized = False
def step(self, closure=None):
loss = None
if closure is not None:
loss = closure()
if not self._momentum_initialized:
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
state = self.state[p]
grad = p.grad.data
if grad.is_sparse:
raise RuntimeError("NovoGrad does not support sparse gradients")
v = torch.norm(grad) ** 2
m = grad / (torch.sqrt(v) + self._eps) + self._wd * p.data
state["step"] = 0
state["v"] = v
state["m"] = m
state["grad_ema"] = None
self._momentum_initialized = True
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
continue
state = self.state[p]
state["step"] += 1
step, v, m = state["step"], state["v"], state["m"]
grad_ema = state["grad_ema"]
grad = p.grad.data
g2 = torch.norm(grad) ** 2
grad_ema = (
g2
if grad_ema is None
else grad_ema * self._beta2 + g2 * (1.0 - self._beta2)
)
grad *= 1.0 / (torch.sqrt(grad_ema) + self._eps)
if self._grad_averaging:
grad *= 1.0 - self._beta1
g2 = torch.norm(grad) ** 2
v = self._beta2 * v + (1.0 - self._beta2) * g2
m = self._beta1 * m + (
grad / (torch.sqrt(v) + self._eps) + self._wd * p.data
)
bias_correction1 = 1 - self._beta1 ** step
bias_correction2 = 1 - self._beta2 ** step
step_size = group["lr"] * math.sqrt(bias_correction2) / bias_correction1
state["v"], state["m"] = v, m
state["grad_ema"] = grad_ema
p.data.add_(-step_size, m)
return loss
|