| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374 |
- r"""Functional interface"""
- import math
- from torch import Tensor
- from typing import List
- from .adadelta import adadelta # type: ignore[attr-defined] # noqa: F401
- from .adagrad import adagrad, _make_sparse # type: ignore[attr-defined] # noqa: F401
- from .adam import adam # type: ignore[attr-defined] # noqa: F401
- from .adamw import adamw # type: ignore[attr-defined] # noqa: F401
- from .adamax import adamax # type: ignore[attr-defined] # noqa: F401
- from .asgd import asgd # type: ignore[attr-defined] # noqa: F401
- from .nadam import nadam # type: ignore[attr-defined] # noqa: F401
- from .radam import radam # type: ignore[attr-defined] # noqa: F401
- from .rmsprop import rmsprop # type: ignore[attr-defined] # noqa: F401
- from .rprop import rprop # type: ignore[attr-defined] # noqa: F401
- from .sgd import sgd # type: ignore[attr-defined] # noqa: F401
- # TODO: use foreach API in optim._functional to do all the computation
- def sparse_adam(params: List[Tensor],
- grads: List[Tensor],
- exp_avgs: List[Tensor],
- exp_avg_sqs: List[Tensor],
- state_steps: List[int],
- *,
- eps: float,
- beta1: float,
- beta2: float,
- lr: float):
- r"""Functional API that performs Sparse Adam algorithm computation.
- See :class:`~torch.optim.SparseAdam` for details.
- """
- for i, param in enumerate(params):
- grad = grads[i]
- grad = grad.coalesce() # the update is non-linear so indices must be unique
- grad_indices = grad._indices()
- grad_values = grad._values()
- size = grad.size()
- exp_avg = exp_avgs[i]
- exp_avg_sq = exp_avg_sqs[i]
- step = state_steps[i]
- def make_sparse(values):
- constructor = grad.new
- if grad_indices.dim() == 0 or values.dim() == 0:
- return constructor().resize_as_(grad)
- return constructor(grad_indices, values, size)
- # Decay the first and second moment running average coefficient
- # old <- b * old + (1 - b) * new
- # <==> old += (1 - b) * (new - old)
- old_exp_avg_values = exp_avg.sparse_mask(grad)._values()
- exp_avg_update_values = grad_values.sub(old_exp_avg_values).mul_(1 - beta1)
- exp_avg.add_(make_sparse(exp_avg_update_values))
- old_exp_avg_sq_values = exp_avg_sq.sparse_mask(grad)._values()
- exp_avg_sq_update_values = grad_values.pow(2).sub_(old_exp_avg_sq_values).mul_(1 - beta2)
- exp_avg_sq.add_(make_sparse(exp_avg_sq_update_values))
- # Dense addition again is intended, avoiding another sparse_mask
- numer = exp_avg_update_values.add_(old_exp_avg_values)
- exp_avg_sq_update_values.add_(old_exp_avg_sq_values)
- denom = exp_avg_sq_update_values.sqrt_().add_(eps)
- del exp_avg_update_values, exp_avg_sq_update_values
- bias_correction1 = 1 - beta1 ** step
- bias_correction2 = 1 - beta2 ** step
- step_size = lr * math.sqrt(bias_correction2) / bias_correction1
- param.add_(make_sparse(-step_size * numer.div_(denom)))
|
