Spaces:
Running
Running
File size: 5,866 Bytes
ab7d699 |
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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import math
import torch
from torch.optim.optimizer import Optimizer
class AdamS(Optimizer):
r"""Implements Adam with stable weight decay (AdamS) algorithm.
It has be proposed in
`Stable Weight Decay Regularization`__.
Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
weight_decay (float, optional): weight decay coefficient (default: 1e-4)
amsgrad (boolean, optional): whether to use the AMSGrad variant of this
algorithm from the paper `On the Convergence of Adam and Beyond`_
(default: False)
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
weight_decay=1e-4, amsgrad=False):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
if not 0.0 <= betas[1] < 1.0:
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
if not 0.0 <= weight_decay:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
defaults = dict(lr=lr, betas=betas, eps=eps,
weight_decay=weight_decay, amsgrad=amsgrad)
super(AdamS, self).__init__(params, defaults)
def __setstate__(self, state):
super(AdamS, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('amsgrad', False)
@torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step.
Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
param_size = 0
exp_avg_sq_hat_sum = 0.
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
param_size += p.numel()
# Perform optimization step
grad = p.grad
if grad.is_sparse:
raise RuntimeError('AdamS does not support sparse gradients')
amsgrad = group['amsgrad']
state = self.state[p]
# State initialization
if len(state) == 0:
state['step'] = 0
# Exponential moving average of gradient values
state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
# Exponential moving average of squared gradient values
state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
if amsgrad:
# Maintains max of all exp. moving avg. of sq. grad. values
state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
beta1, beta2 = group['betas']
exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
state['step'] += 1
bias_correction2 = 1 - beta2 ** state['step']
# Decay the first and second moment running average coefficient
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
if amsgrad:
max_exp_avg_sq = state['max_exp_avg_sq']
# Maintains the maximum of all 2nd moment running avg. till now
torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
# Use the max. for normalizing running avg. of gradient
exp_avg_sq_hat = max_exp_avg_sq / bias_correction2
else:
exp_avg_sq_hat = exp_avg_sq / bias_correction2
exp_avg_sq_hat_sum += exp_avg_sq_hat.sum()
# Calculate the sqrt of the mean of all elements in exp_avg_sq_hat
exp_avg_mean_sqrt = math.sqrt(exp_avg_sq_hat_sum / param_size)
for group in self.param_groups:
for p in group['params']:
if p.grad is None:
continue
state = self.state[p]
#Perform stable weight decay
if group['weight_decay'] !=0:
p.data.mul_(1 - group['weight_decay'] * group['lr'] / exp_avg_mean_sqrt)
beta1, beta2 = group['betas']
exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
bias_correction1 = 1 - beta1 ** state['step']
bias_correction2 = 1 - beta2 ** state['step']
if amsgrad:
max_exp_avg_sq = state['max_exp_avg_sq']
exp_avg_sq_hat = max_exp_avg_sq / bias_correction2
else:
exp_avg_sq_hat = exp_avg_sq / bias_correction2
denom = exp_avg_sq_hat.sqrt().add(group['eps'])
step_size = group['lr'] / bias_correction1
p.addcdiv_(exp_avg, denom, value= - step_size)
return loss
|