def adadelta(params: List[Tensor], grads: List[Tensor],
square_avgs: List[Tensor], acc_deltas: List[Tensor], *, lr: float,
weight_decay: float, rho: float, eps: float):
r"""Functional API that performs Adadelta algorithm computation.
See :class:`~torch.optim.Adadelta` for details.
"""
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
torch._foreach_mul_(square_avgs, rho)
torch._foreach_addcmul_(square_avgs, grads, grads, value=1 - rho)
std = torch._foreach_add(square_avgs, eps)
torch._foreach_sqrt_(std)
deltas = torch._foreach_add(acc_deltas, eps)
torch._foreach_sqrt_(deltas)
torch._foreach_div_(deltas, std)
torch._foreach_mul_(deltas, grads)
torch._foreach_add_(params, deltas, alpha=-lr)
torch._foreach_mul_(acc_deltas, rho)
torch._foreach_addcmul_(acc_deltas, deltas, deltas, value=1 - rho)
def asgd(params: List[Tensor], grads: List[Tensor], states: List[Dict],
lambd: float, lr: float, t0: float, alpha: float,
weight_decay: float):
r"""Functional API that performs ASGD algorithm computation.
See :class:`~torch.optim.ASGD` for details.
"""
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
# decay term
eta = states[0]['eta']
torch._foreach_mul_(params, 1 - lambd * eta)
# update parameter
torch._foreach_add_(params, grads, alpha=-eta)
# averaging
for i in range(len(states)):
if states[i]['mu'] != 1:
states[i]['ax'].add_(params[i].sub(states[i]['ax']).mul(
states[i]['mu']))
else:
states[i]['ax'].copy_(params[i])
# update eta and mu
for state in states:
state['eta'] = (lr / math.pow((1 + lambd * lr * state['step']), alpha))
state['mu'] = 1 / max(1, state['step'] - t0)
def _multi_tensor_rmsprop(params: List[Tensor], grads: List[Tensor],
square_avgs: List[Tensor], grad_avgs: List[Tensor],
momentum_buffer_list: List[Tensor], *, lr: float,
alpha: float, eps: float, weight_decay: float,
momentum: float, centered: bool):
if len(params) == 0:
return
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
torch._foreach_mul_(square_avgs, alpha)
torch._foreach_addcmul_(square_avgs, grads, grads, value=1 - alpha)
if centered:
torch._foreach_mul_(grad_avgs, alpha)
torch._foreach_add_(grad_avgs, grads, alpha=1 - alpha)
avg = torch._foreach_addcmul(square_avgs,
grad_avgs,
grad_avgs,
value=-1)
torch._foreach_sqrt_(avg)
torch._foreach_add_(avg, eps)
else:
avg = torch._foreach_sqrt(square_avgs)
torch._foreach_add_(avg, eps)
if momentum > 0:
torch._foreach_mul_(momentum_buffer_list, momentum)
torch._foreach_addcdiv_(momentum_buffer_list, grads, avg)
torch._foreach_add_(params, momentum_buffer_list, alpha=-lr)
else:
torch._foreach_addcdiv_(params, grads, avg, value=-lr)
def _multi_tensor_adadelta(params: List[Tensor],
grads: List[Tensor],
square_avgs: List[Tensor],
acc_deltas: List[Tensor],
*,
lr: float,
weight_decay: float,
rho: float,
eps: float):
if len(params) == 0:
return
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
torch._foreach_mul_(square_avgs, rho)
torch._foreach_addcmul_(square_avgs, grads, grads, value=1 - rho)
std = torch._foreach_add(square_avgs, eps)
torch._foreach_sqrt_(std)
deltas = torch._foreach_add(acc_deltas, eps)
torch._foreach_sqrt_(deltas)
torch._foreach_div_(deltas, std)
torch._foreach_mul_(deltas, grads)
torch._foreach_add_(params, deltas, alpha=-lr)
torch._foreach_mul_(acc_deltas, rho)
torch._foreach_addcmul_(acc_deltas, deltas, deltas, value=1 - rho)
def test_complex_scalar(self, device, dtype):
tensors = [
torch.zeros(10, 10, device=device, dtype=dtype) for _ in range(10)
]
complex_scalar = 3 + 5j
# bool tensor + 1 will result in int64 tensor
expected = [
torch.add(complex_scalar,
torch.zeros(10, 10, device=device, dtype=dtype))
for _ in range(10)
]
if dtype in [
torch.float16, torch.float32, torch.float64, torch.bfloat16
] and device == 'cuda:0':
# value cannot be converted to dtype without overflow:
self.assertRaises(
RuntimeError,
lambda: torch._foreach_add_(tensors, complex_scalar))
self.assertRaises(
RuntimeError,
lambda: torch._foreach_add(tensors, complex_scalar))
return
res = torch._foreach_add(tensors, complex_scalar)
self.assertEqual(res, expected)
if dtype not in [torch.complex64, torch.complex128]:
self.assertRaises(
RuntimeError,
lambda: torch._foreach_add_(tensors, complex_scalar))
else:
torch._foreach_add_(tensors, complex_scalar)
self.assertEqual(res, tensors)
def _update(self):
# _foreach_** is n times faster than for loops
o_p = [
p.data for p in self._original_model.parameters()
if isinstance(p, torch.Tensor)
]
e_p = [
p.data for p in self._ema_model.parameters()
if isinstance(p, torch.Tensor)
]
torch._foreach_mul_(e_p, self.momentum)
torch._foreach_add_(e_p, o_p, alpha=1 - self.momentum)
# some buffers are integer for counting etc.
o_b = [
b for b in self._original_model.buffers()
if isinstance(b, torch.Tensor) and torch.is_floating_point(b)
]
if len(o_b) > 0:
e_b = [
b for b in self._ema_model.buffers()
if isinstance(b, torch.Tensor) and torch.is_floating_point(b)
]
torch._foreach_mul_(e_b, self.momentum)
torch._foreach_add_(e_b, o_b, alpha=1 - self.momentum)
def test_int_scalar(self, device, dtype):
tensors = [
torch.zeros(10, 10, device=device, dtype=dtype) for _ in range(10)
]
int_scalar = 1
# bool tensor + 1 will result in int64 tensor
if dtype == torch.bool:
expected = [
torch.ones(10, 10, device=device, dtype=torch.int64)
for _ in range(10)
]
else:
expected = [
torch.ones(10, 10, device=device, dtype=dtype)
for _ in range(10)
]
res = torch._foreach_add(tensors, int_scalar)
self.assertEqual(res, expected)
if dtype in [torch.bool]:
with self.assertRaisesRegex(
RuntimeError,
"result type Long can't be cast to the desired output type Bool"
):
torch._foreach_add_(tensors, int_scalar)
else:
torch._foreach_add_(tensors, int_scalar)
self.assertEqual(res, tensors)
def test_float_scalar(self, device, dtype):
tensors = [
torch.zeros(10, 10, device=device, dtype=dtype) for _ in range(10)
]
float_scalar = 1.
# float scalar + integral tensor will result in float tensor
if dtype in [
torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64,
torch.bool
]:
expected = [
torch.ones(10, 10, device=device, dtype=torch.float32)
for _ in range(10)
]
else:
expected = [
torch.ones(10, 10, device=device, dtype=dtype)
for _ in range(10)
]
res = torch._foreach_add(tensors, float_scalar)
self.assertEqual(res, expected)
if dtype in [
torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64,
torch.bool
]:
self.assertRaises(
RuntimeError,
lambda: torch._foreach_add_(tensors, float_scalar))
else:
torch._foreach_add_(tensors, float_scalar)
self.assertEqual(res, tensors)
def adamax(params: List[Tensor], grads: List[Tensor], exp_avgs: List[Tensor],
exp_infs: List[Tensor], states: List[Dict], *, beta1: float,
beta2: float, lr: float, weight_decay: float, eps: float):
r"""Functional API that performs Adamax algorithm computation.
See :class:`~torch.optim.Adamax` for details.
"""
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
# Update biased first moment estimate.
torch._foreach_mul_(exp_avgs, beta1)
torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)
# Update the exponentially weighted infinity norm.
torch._foreach_mul_(exp_infs, beta2)
for exp_inf, grad in zip(exp_infs, grads):
norm_buf = torch.cat(
[exp_inf.unsqueeze(0),
grad.abs().add_(eps).unsqueeze_(0)], 0)
torch.max(norm_buf,
0,
keepdim=False,
out=(exp_inf, exp_inf.new().long()))
bias_corrections = [1 - beta1**state['step'] for state in states]
clr = [-1 * (lr / bias_correction) for bias_correction in bias_corrections]
torch._foreach_addcdiv_(params, exp_avgs, exp_infs, clr)
def test_add_with_different_size_tensors(self, device, dtype):
if dtype == torch.bool:
return
tensors = [torch.zeros(10 + n, 10 + n, device=device, dtype=dtype) for n in range(10)]
expected = [torch.ones(10 + n, 10 + n, device=device, dtype=dtype) for n in range(10)]
torch._foreach_add_(tensors, 1)
self.assertEqual(expected, tensors)
def test_add_list_different_sizes(self, device, dtype):
tensors1 = [torch.zeros(10 + n, 10 + n, device=device, dtype=dtype) for n in range(10)]
tensors2 = [torch.ones(10 + n, 10 + n, device=device, dtype=dtype) for n in range(10)]
res = torch._foreach_add(tensors1, tensors2)
torch._foreach_add_(tensors1, tensors2)
self.assertEqual(res, tensors1)
self.assertEqual(res, [torch.ones(10 + n, 10 + n, device=device, dtype=dtype) for n in range(10)])
def test_add_scalar_with_empty_list_and_empty_tensor(self, device, dtype):
# TODO: enable empty list case
for tensors in [[torch.randn([0])]]:
res = torch._foreach_add(tensors, 1)
self.assertEqual(res, tensors)
torch._foreach_add_(tensors, 1)
self.assertEqual(res, tensors)
def adagrad(params: List[Tensor], grads: List[Tensor],
state_sums: List[Tensor], state_steps: List[int],
has_sparse_grad: bool, *, lr: float, weight_decay: float,
lr_decay: float, eps: float):
r"""Functional API that performs Adagrad algorithm computation.
See :class:`~torch.optim.Adagrad` for details.
"""
if weight_decay != 0:
if has_sparse_grad:
raise RuntimeError(
"weight_decay option is not compatible with sparse gradients")
torch._foreach_add_(grads, params, alpha=weight_decay)
minus_clr = [-lr / (1 + (step - 1) * lr_decay) for step in state_steps]
if has_sparse_grad:
# sparse is not supported by multi_tensor. Fall back to optim.adagrad
# implementation for sparse gradients
for i, (param, grad, state_sum,
step) in enumerate(zip(params, grads, state_sums,
state_steps)):
grad = grad.coalesce(
) # the update is non-linear so indices must be unique
grad_indices = grad._indices()
grad_values = grad._values()
size = grad.size()
state_sum.add_(_make_sparse(grad, grad_indices,
grad_values.pow(2)))
std_sparse = state_sum.sparse_mask(grad)
std_sparse_values = std_sparse._values().sqrt_().add_(eps)
param.add_(
_make_sparse(grad, grad_indices,
grad_values / std_sparse_values),
alpha=minus_clr[i],
)
else:
grads = [
torch.view_as_real(x) if torch.is_complex(x) else x for x in grads
]
state_sums = [
torch.view_as_real(x) if torch.is_complex(x) else x
for x in state_sums
]
torch._foreach_addcmul_(state_sums, grads, grads, value=1)
std = torch._foreach_add(torch._foreach_sqrt(state_sums), eps)
toAdd = torch._foreach_div(torch._foreach_mul(grads, minus_clr), std)
toAdd = [
torch.view_as_complex(x) if torch.is_complex(params[i]) else x
for i, x in enumerate(toAdd)
]
torch._foreach_add_(params, toAdd)
state_sums = [
torch.view_as_complex(x) if torch.is_complex(params[i]) else x
for i, x in enumerate(state_sums)
]
def pre_step(self):
if self.pre_op:
with torch.no_grad():
params, grads = zip(*self.params_with_grads())
torch._foreach_add_(grads, params, alpha=self.value)
if self.log:
logging.debug(
'L2 penalty of %s was applied pre optimization step',
self.value)
def test_bool_scalar(self, device, dtype):
tensors = [torch.zeros(10, 10, device=device, dtype=dtype) for _ in range(10)]
bool_scalar = True
expected = [torch.ones(10, 10, device=device, dtype=dtype) for _ in range(10)]
res = torch._foreach_add(tensors, bool_scalar)
self.assertEqual(res, expected)
torch._foreach_add_(tensors, bool_scalar)
self.assertEqual(res, tensors)
def _multi_tensor_rmsprop(params: List[Tensor], grads: List[Tensor],
square_avgs: List[Tensor], grad_avgs: List[Tensor],
momentum_buffer_list: List[Tensor], *, lr: float,
alpha: float, eps: float, weight_decay: float,
momentum: float, centered: bool, maximize: bool,
differentiable: bool):
if len(params) == 0:
return
assert not differentiable, "_foreach ops don't support autograd"
if maximize:
grads = torch._foreach_neg(grads)
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
def _view_complex_as_real(tensor_list):
return [
torch.view_as_real(t) if torch.is_complex(t) else t
for t in tensor_list
]
grads = _view_complex_as_real(grads)
params = _view_complex_as_real(params)
square_avgs = _view_complex_as_real(square_avgs)
torch._foreach_mul_(square_avgs, alpha)
torch._foreach_addcmul_(square_avgs, grads, grads, value=1 - alpha)
if centered:
grad_avgs = _view_complex_as_real(grad_avgs)
torch._foreach_mul_(grad_avgs, alpha)
torch._foreach_add_(grad_avgs, grads, alpha=1 - alpha)
avg = torch._foreach_addcmul(square_avgs,
grad_avgs,
grad_avgs,
value=-1)
torch._foreach_sqrt_(avg)
torch._foreach_add_(avg, eps)
else:
avg = torch._foreach_sqrt(square_avgs)
torch._foreach_add_(avg, eps)
if momentum > 0:
momentum_buffer_list = _view_complex_as_real(momentum_buffer_list)
torch._foreach_mul_(momentum_buffer_list, momentum)
torch._foreach_addcdiv_(momentum_buffer_list, grads, avg)
torch._foreach_add_(params, momentum_buffer_list, alpha=-lr)
else:
torch._foreach_addcdiv_(params, grads, avg, value=-lr)
def test_add_list_same_size(self, device, dtype):
tensors1 = [
torch.zeros(10, 10, device=device, dtype=dtype) for _ in range(10)
]
tensors2 = [
torch.ones(10, 10, device=device, dtype=dtype) for _ in range(10)
]
res = torch._foreach_add(tensors1, tensors2)
torch._foreach_add_(tensors1, tensors2)
self.assertEqual(res, tensors1)
self.assertEqual(res[0], torch.ones(10, 10, device=device,
dtype=dtype))
def radam(params: List[Tensor], grads: List[Tensor], exp_avg: List[Tensor],
exp_avg_sq: List[Tensor], states: List[Dict], *, beta1: float,
beta2: float, lr: float, weight_decay: float, eps: float):
r"""Functional API that performs RAdam algorithm computation.
See :class:`~torch.optim.RAdam` for details.
"""
# maximum length of the approximated SMA
rho_inf = 2 / (1 - beta2) - 1
# compute the length of the approximated SMA
rho_t_list = [
rho_inf - 2 * state['step'] * (beta2**state['step']) /
(1 - beta2**state['step']) for state in states
]
bias_correction1 = [1 - beta1**state['step'] for state in states]
bias_correction2 = [1 - beta2**state['step'] for state in states]
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
# Decay the first and second moment running average coefficient
torch._foreach_mul_(exp_avg, beta1)
torch._foreach_add_(exp_avg, grads, alpha=1 - beta1)
torch._foreach_mul_(exp_avg_sq, beta2)
torch._foreach_addcmul_(exp_avg_sq, grads, grads, 1 - beta2)
rect = [
math.sqrt((rho_t - 4) * (rho_t - 2) * rho_inf /
((rho_inf - 4) * (rho_inf - 2) * rho_t)) if rho_t > 5 else 0
for rho_t in rho_t_list
]
unrectified = [0 if rect > 0 else 1. for rect in rect]
exp_avg_sq_sqrt = torch._foreach_sqrt(exp_avg_sq)
bias_correction_sqrt = [math.sqrt(bc) for bc in bias_correction2]
denom = torch._foreach_div(exp_avg_sq_sqrt, bias_correction_sqrt)
step_size = [(lr * rect / bc) * -1
for rect, bc in zip(rect, bias_correction1)]
torch._foreach_addcdiv_(params, exp_avg, denom, step_size)
denom = [
torch.ones_like(exp_av, memory_format=torch.preserve_format)
for exp_av in exp_avg
]
step_size = [(lr * rect / bc) * -1
for rect, bc in zip(unrectified, bias_correction1)]
torch._foreach_addcdiv_(params, exp_avg, denom, step_size)
def update(self, model):
x = []
y = []
needs_module = hasattr(model, 'module') and not self.ema_has_module
with torch.no_grad():
for ema_v, model_v in zip(self.ema.state_dict().values(),
model.state_dict().values()):
x.append(ema_v.type(torch.float32))
if self.device:
model_v = model_v.detach().to(device=self.device)
y.append(model_v.type(torch.float32))
torch._foreach_mul_(x, self.decay)
torch._foreach_add_(x, y, alpha=1. - self.decay)
for ind, ema_v in enumerate(self.ema.state_dict().values()):
ema_v.copy_(x[ind])
def test_add_list_slow_path(self, device, dtype):
# different strides
tensor1 = torch.zeros(10, 10, device=device, dtype=dtype)
tensor2 = torch.ones(10, 10, device=device, dtype=dtype)
res = torch._foreach_add([tensor1], [tensor2.t()])
torch._foreach_add_([tensor1], [tensor2])
self.assertEqual(res, [tensor1])
# non contiguous
tensor1 = torch.randn(5, 2, 1, 3, device=device)[:, 0]
tensor2 = torch.randn(5, 2, 1, 3, device=device)[:, 0]
self.assertFalse(tensor1.is_contiguous())
self.assertFalse(tensor2.is_contiguous())
res = torch._foreach_add([tensor1], [tensor2])
torch._foreach_add_([tensor1], [tensor2])
self.assertEqual(res, [tensor1])
def _multi_tensor_adamw(params: List[Tensor], grads: List[Tensor],
exp_avgs: List[Tensor], exp_avg_sqs: List[Tensor],
max_exp_avg_sqs: List[Tensor],
state_steps: List[Tensor], *, amsgrad: bool,
beta1: float, beta2: float, lr: float,
weight_decay: float, eps: float, maximize: bool):
if len(params) == 0:
return
if maximize:
grads = torch._foreach_neg(tuple(grads)) # type: ignore[assignment]
# Perform stepweight decay
torch._foreach_mul_(params, 1 - lr * weight_decay)
# update steps
torch._foreach_add_(state_steps, 1)
bias_correction1 = [1 - beta1**step.item() for step in state_steps]
bias_correction2 = [1 - beta2**step.item() for step in state_steps]
# Decay the first and second moment running average coefficient
torch._foreach_mul_(exp_avgs, beta1)
torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)
torch._foreach_mul_(exp_avg_sqs, beta2)
torch._foreach_addcmul_(exp_avg_sqs, grads, grads, 1 - beta2)
if amsgrad:
# Maintains the maximum of all 2nd moment running avg. till now
max_exp_avg_sqs = torch._foreach_maximum(
max_exp_avg_sqs, exp_avg_sqs) # type: ignore[assignment]
# Use the max. for normalizing running avg. of gradient
max_exp_avg_sq_sqrt = torch._foreach_sqrt(max_exp_avg_sqs)
bias_correction_sqrt = [math.sqrt(bc) for bc in bias_correction2]
torch._foreach_div_(max_exp_avg_sq_sqrt, bias_correction_sqrt)
denom = torch._foreach_add(max_exp_avg_sq_sqrt, eps)
else:
exp_avg_sq_sqrt = torch._foreach_sqrt(exp_avg_sqs)
bias_correction_sqrt = [math.sqrt(bc) for bc in bias_correction2]
torch._foreach_div_(exp_avg_sq_sqrt, bias_correction_sqrt)
denom = torch._foreach_add(exp_avg_sq_sqrt, eps)
step_size = [-1 * (lr / bc) for bc in bias_correction1]
torch._foreach_addcdiv_(params, exp_avgs, denom, step_size)
def on_step(self, task) -> None:
if not task.train:
return
with torch.no_grad():
if self.use_optimization(task):
torch._foreach_mul_(self.ema_model_state_list, self.decay)
torch._foreach_add_(self.ema_model_state_list,
self.param_list,
alpha=(1 - self.decay))
else:
for name, param in self.get_model_state_iterator(
task.base_model):
self.state.ema_model_state[
name] = self.decay * self.state.ema_model_state[
name] + (1 -
self.decay) * param.to(device=self.device)
def multi_tensor_scale(
self,
src: Sequence[torch.Tensor],
dst: Sequence[torch.Tensor],
scale: float,
) -> None:
with torch.no_grad(): # type: ignore[no-untyped-call]
# _foreach_zero for long type is not supported in CUDA
if self._enable_foreach and src[0].is_floating_point():
# scale
val = torch._foreach_mul(tuple(src), scale)
# copy tensor
torch._foreach_zero_(tuple(dst))
torch._foreach_add_(tuple(dst), val)
else:
for s, d in zip(src, dst):
d.copy_(s * scale)
def nadam(params: List[Tensor],
grads: List[Tensor],
exp_avg: List[Tensor],
exp_avg_sq: List[Tensor],
mu_products: List[Tensor],
states: List[Dict],
*,
beta1: float,
beta2: float,
lr: float,
weight_decay: float,
momentum_decay: float,
eps: float):
r"""Functional API that performs NAdam algorithm computation.
See :class:`~torch.optim.NAdam` for details.
"""
bias_correction1 = [1 - beta1 ** state['step'] for state in states]
bias_correction2 = [1 - beta2 ** state['step'] for state in states]
mus = [beta1 * (1. - 0.5 * (0.96 ** (state['step'] * momentum_decay))) for state in states]
mu_nexts = [beta1 * (1. - 0.5 * (0.96 ** ((state['step'] + 1) * momentum_decay)))
for state in states]
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
# Decay the first and second moment running average coefficient
torch._foreach_mul_(exp_avg, beta1)
torch._foreach_add_(exp_avg, grads, alpha=1 - beta1)
torch._foreach_mul_(exp_avg_sq, beta2)
torch._foreach_addcmul_(exp_avg_sq, grads, grads, 1 - beta2)
exp_avg_sq_sqrt = torch._foreach_sqrt(exp_avg_sq)
bias_correction_sqrt = [math.sqrt(bc) for bc in bias_correction2]
torch._foreach_div_(exp_avg_sq_sqrt, bias_correction_sqrt)
denom = torch._foreach_add(exp_avg_sq_sqrt, eps)
step_size_grads = [(lr * (1. - mu) / (1. - mu_product)) * -1
for mu_product, mu in zip(mu_products, mus)]
step_size_expavg = [(lr * mu_next / (1. - mu_product * mu_next)) * -1
for mu_product, mu_next in zip(mu_products, mu_nexts)]
torch._foreach_addcdiv_(params, grads, denom, step_size_grads)
torch._foreach_addcdiv_(params, exp_avg, denom, step_size_expavg)
def step(self):
weight_decays = []
for group in self.optim.param_groups:
# absorb weight decay control from optimizer
weight_decay = group[
'weight_decay'] if 'weight_decay' in group else 0
weight_decays.append(weight_decay)
group['weight_decay'] = 0
params = []
grads = []
lrs = []
for p in group['params']:
if p.grad is None:
continue
param_norm = torch.norm(p.data)
grad_norm = torch.norm(p.grad.data)
if param_norm != 0 and grad_norm != 0:
# calculate adaptive lr + weight decay
# .item() may be sub-optimal, but required because _foreach_* don't support broadcasting at the moment
adaptive_lr = (self.trust_coefficient * param_norm /
(grad_norm + param_norm * weight_decay +
self.eps)).item()
# clip learning rate for LARC
if self.clip:
# calculation of adaptive_lr so that when multiplied by lr it equals `min(adaptive_lr, lr)`
adaptive_lr = min(adaptive_lr / group['lr'], 1.0)
params.append(p.data)
grads.append(p.grad.data)
lrs.append(adaptive_lr)
# p.grad.data += weight_decay * p.data
# p.grad.data *= adaptive_lr
torch._foreach_add_(grads, params, alpha=weight_decay)
torch._foreach_mul_(grads, lrs)
self.optim.step()
# return weight decay control to optimizer
for i, group in enumerate(self.optim.param_groups):
group['weight_decay'] = weight_decays[i]
def step(self,
closure
) -> torch.Tensor:
"""
Args:
closure: A closure that reevaluates the model and returns the loss.
Returns: the loss value evaluated on the original point
"""
closure = torch.enable_grad()(closure)
loss = closure().detach()
for group in self.param_groups:
grads = []
params_with_grads = []
rho = group['rho']
# update internal_optim's learning rate
for p in group['params']:
if p.grad is not None:
# without clone().detach(), p.grad will be zeroed by closure()
grads.append(p.grad.clone().detach())
params_with_grads.append(p)
device = grads[0].device
# compute \hat{\epsilon}=\rho/\norm{g}\|g\|
grad_norm = torch.stack([g.detach().norm(2).to(device) for g in grads]).norm(2)
epsilon = grads # alias for readability
torch._foreach_mul_(epsilon, rho / grad_norm)
# virtual step toward \epsilon
torch._foreach_add_(params_with_grads, epsilon)
# compute g=\nabla_w L_B(w)|_{w+\hat{\epsilon}}
closure()
# virtual step back to the original point
torch._foreach_sub_(params_with_grads, epsilon)
super().step()
return loss
def _update(self):
if torch.cuda.is_available():
torch.cuda.synchronize()
# _foreach_** is n times faster than for loops
o_p = [
p.data for p in self._original_model.parameters()
if isinstance(p, torch.Tensor)
]
e_p = [
p.data for p in self._ema_model.parameters()
if isinstance(p, torch.Tensor)
]
torch._foreach_mul_(e_p, self.momentum)
torch._foreach_add_(e_p, o_p, alpha=1 - self.momentum)
# some buffers are integer for counting etc.
alpha = 0 if self.copy_buffer else self.momentum
o_b = [
b for b in self._original_model.buffers()
if isinstance(b, torch.Tensor) and torch.is_floating_point(b)
]
if len(o_b) > 0:
e_b = [
b for b in self._ema_model.buffers()
if isinstance(b, torch.Tensor) and torch.is_floating_point(b)
]
torch._foreach_mul_(e_b, alpha)
torch._foreach_add_(e_b, o_b, alpha=1 - alpha)
# integers
o_b = [
b for b in self._original_model.buffers()
if isinstance(b, torch.Tensor) and not torch.is_floating_point(b)
]
if len(o_b) > 0:
e_b = [
b for b in self._ema_model.buffers() if
isinstance(b, torch.Tensor) and not torch.is_floating_point(b)
]
for o, e in zip(o_b, e_b):
e.copy_(o)
def _multi_tensor_sgd(params: List[Tensor],
grads: List[Tensor],
momentum_buffer_list: List[Optional[Tensor]],
*,
weight_decay: float,
momentum: float,
lr: float,
dampening: float,
nesterov: bool,
maximize: bool,
has_sparse_grad: bool):
if len(params) == 0:
return
if has_sparse_grad is None:
has_sparse_grad = any(grad.is_sparse for grad in grads)
if maximize:
grads = torch._foreach_neg(tuple(grads)) # type: ignore[assignment]
if weight_decay != 0:
grads = torch._foreach_add(grads, params, alpha=weight_decay)
if momentum != 0:
bufs = []
all_states_with_momentum_buffer = True
for i in range(len(momentum_buffer_list)):
if momentum_buffer_list[i] is None:
all_states_with_momentum_buffer = False
break
else:
bufs.append(momentum_buffer_list[i])
if all_states_with_momentum_buffer:
torch._foreach_mul_(bufs, momentum)
torch._foreach_add_(bufs, grads, alpha=1 - dampening)
else:
bufs = []
for i in range(len(momentum_buffer_list)):
if momentum_buffer_list[i] is None:
buf = momentum_buffer_list[i] = torch.clone(grads[i]).detach()
else:
buf = momentum_buffer_list[i]
buf.mul_(momentum).add_(grads[i], alpha=1 - dampening)
bufs.append(buf)
if nesterov:
torch._foreach_add_(grads, bufs, alpha=momentum)
else:
grads = bufs
if not has_sparse_grad:
torch._foreach_add_(params, grads, alpha=-lr)
else:
# foreach APIs dont support sparse
for i in range(len(params)):
params[i].add_(grads[i], alpha=-lr)
def _multi_tensor_adamax(params: List[Tensor], grads: List[Tensor],
exp_avgs: List[Tensor], exp_infs: List[Tensor],
state_steps: List[Tensor], *, beta1: float,
beta2: float, lr: float, weight_decay: float,
eps: float):
if len(params) == 0:
return
# Update steps
torch._foreach_add_(state_steps, 1)
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
# Update biased first moment estimate.
torch._foreach_mul_(exp_avgs, beta1)
torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)
# Update the exponentially weighted infinity norm.
torch._foreach_mul_(exp_infs, beta2)
for exp_inf, grad in zip(exp_infs, grads):
norm_buf = torch.cat(
[exp_inf.unsqueeze(0),
grad.abs().add_(eps).unsqueeze_(0)], 0)
torch.max(norm_buf,
0,
keepdim=False,
out=(exp_inf, exp_inf.new().long()))
bias_corrections = [1 - beta1**step.item() for step in state_steps]
clr = [-1 * (lr / bias_correction) for bias_correction in bias_corrections]
torch._foreach_addcdiv_(params, exp_avgs, exp_infs, clr)
def _multi_tensor_radam(params: List[Tensor], grads: List[Tensor],
exp_avgs: List[Tensor], exp_avg_sqs: List[Tensor],
state_steps: List[Tensor], *, beta1: float,
beta2: float, lr: float, weight_decay: float,
eps: float):
if len(params) == 0:
return
# Update steps
torch._foreach_add_(state_steps, 1)
# maximum length of the approximated SMA
rho_inf = 2 / (1 - beta2) - 1
# compute the length of the approximated SMA
rho_t_list = [
rho_inf - 2 * step.item() * (beta2**step.item()) /
(1 - beta2**step.item()) for step in state_steps
]
bias_correction1 = [1 - beta1**step.item() for step in state_steps]
bias_correction2 = [1 - beta2**step.item() for step in state_steps]
if weight_decay != 0:
torch._foreach_add_(grads, params, alpha=weight_decay)
# Decay the first and second moment running average coefficient
torch._foreach_mul_(exp_avgs, beta1)
torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)
torch._foreach_mul_(exp_avg_sqs, beta2)
torch._foreach_addcmul_(exp_avg_sqs, grads, grads, 1 - beta2)
rect = [
math.sqrt((rho_t - 4) * (rho_t - 2) * rho_inf /
((rho_inf - 4) * (rho_inf - 2) * rho_t)) if rho_t > 5 else 0
for rho_t in rho_t_list
]
unrectified = [0 if rect > 0 else 1. for rect in rect]
exp_avg_sq_sqrt = torch._foreach_sqrt(exp_avg_sqs)
bias_correction_sqrt = [math.sqrt(bc) for bc in bias_correction2]
denom = torch._foreach_div(exp_avg_sq_sqrt, bias_correction_sqrt)
step_size = [(lr * rect / bc) * -1
for rect, bc in zip(rect, bias_correction1)]
torch._foreach_addcdiv_(params, exp_avgs, denom, step_size)
denom = [
torch.ones_like(exp_av, memory_format=torch.preserve_format)
for exp_av in exp_avgs
]
step_size = [(lr * rect / bc) * -1
for rect, bc in zip(unrectified, bias_correction1)]
torch._foreach_addcdiv_(params, exp_avgs, denom, step_size)
