Skip to content

[Optim] Upgrade adam to 1.9 #39

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
eric-zheng wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Optim] Upgrade adam to 1.9 #39

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 24 additions & 11 deletions hfta/optim/adam.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,10 @@ class Adam(Optimizer):
r"""Implements Adam algorithm.

It has been proposed in `Adam: A Method for Stochastic Optimization`_.
The implementation of the L2 penalty follows changes proposed in
`Decoupled Weight Decay Regularization`_.

Arguments:
Args:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float or a list/tuple/np.array/torch.Tensor of floats, optional):
Expand All @@ -32,6 +34,8 @@ class Adam(Optimizer):

.. _Adam\: A Method for Stochastic Optimization:
https://arxiv.org/abs/1412.6980
.. _Decoupled Weight Decay Regularization:
https://arxiv.org/abs/1711.05101
.. _On the Convergence of Adam and Beyond:
https://openreview.net/forum?id=ryQu7f-RZ
"""
Expand Down Expand Up @@ -81,7 +85,7 @@ def zero_grad(self):
def step(self, closure=None):
"""Performs a single optimization step.

Arguments:
Args:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
Expand All @@ -91,6 +95,8 @@ def step(self, closure=None):
loss = closure()

for group in self.param_groups:
beta1, beta2 = group['betas']

for p in group['params']:
if p.grad is None:
continue
Expand Down Expand Up @@ -123,18 +129,22 @@ def step(self, closure=None):
exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
if amsgrad:
max_exp_avg_sq = state['max_exp_avg_sq']
beta1, beta2 = group['betas']
lr, eps, weight_decay = group['lr'], group['eps'], group['weight_decay']

state['step'] += 1
if isinstance(beta1, dict):
bias_correction1 = 1 - beta1[p]**state['step']
else:
bias_correction1 = 1 - beta1**state['step']
if isinstance(beta2, dict):
sqrt_bias_correction2 = (1 - beta2[p]**state['step']).sqrt()
else:
sqrt_bias_correction2 = math.sqrt(1 - beta2**state['step'])

# Start of functional computation
def get_bias_correction(beta, step):
if isinstance(beta, dict):
bias_correction = 1 - beta[p]**step
else:
bias_correction = 1 - beta**step
return bias_correction

bias_correction1 = get_bias_correction(beta1, state['step'])
bias_correction2 = get_bias_correction(beta2, state['step'])
sqrt_bias_correction2 = bias_correction2.sqrt() if isinstance(
bias_correction2, torch.Tensor) else math.sqrt(bias_correction2)

if isinstance(weight_decay, dict) or weight_decay != 0:
if isinstance(weight_decay, dict):
Expand All @@ -147,10 +157,12 @@ def step(self, closure=None):
exp_avg.mul_(beta1[p]).add_((1 - beta1[p]) * grad)
else:
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)

if isinstance(beta2, dict):
exp_avg_sq.mul_(beta2[p]).add_((1 - beta2[p]) * grad * grad)
else:
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

if amsgrad:
# 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)
Expand All @@ -169,6 +181,7 @@ def step(self, closure=None):
step_size = lr[p] / bias_correction1
else:
step_size = lr / bias_correction1

if isinstance(step_size, torch.Tensor):
p.add_(-step_size * (exp_avg / denom))
else:
Expand Down