File size: 3,399 Bytes
c61ccee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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,

                maximize: bool):
    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 if not maximize else -grad
        grad = grad.coalesce()  # the update is non-linear so indices must be unique
        grad_indices = grad._indices()
        grad_values = grad._values()
        if grad_values.numel() == 0:
            # Skip update for empty grad
            continue
        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)))