RepDL/repdl/optim.py at main · microsoft/RepDL · GitHub

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
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import math

import torch

from . import ops


class Adam(torch.optim.Adam):
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        for group in self.param_groups:
            params_with_grad = []
            grads = []
            exp_avgs = []
            exp_avg_sqs = []
            max_exp_avg_sqs = []
            state_steps = []
            beta1, beta2 = group["betas"]
            for p in group["params"]:
                if p.grad is not None:
                    params_with_grad.append(p)
                    if p.grad.is_sparse:
                        raise RuntimeError(
                            "Adam does not support sparse gradients, please consider SparseAdam instead"
                        )
                    grads.append(p.grad)
                    state = self.state[p]
                    if len(state) == 0:
                        state["step"] = 0
                        state["exp_avg"] = torch.zeros_like(
                            p, memory_format=torch.preserve_format
                        )
                        state["exp_avg_sq"] = torch.zeros_like(
                            p, memory_format=torch.preserve_format
                        )
                        if group["amsgrad"]:
                            state["max_exp_avg_sq"] = torch.zeros_like(
                                p, memory_format=torch.preserve_format
                            )
                    exp_avgs.append(state["exp_avg"])
                    exp_avg_sqs.append(state["exp_avg_sq"])
                    if group["amsgrad"]:
                        max_exp_avg_sqs.append(state["max_exp_avg_sq"])
                    state["step"] += 1
                    state_steps.append(state["step"])

            # Changed:

            for i, param in enumerate(params_with_grad):
                grad = grads[i]
                exp_avg = exp_avgs[i]
                exp_avg_sq = exp_avg_sqs[i]
                step = state_steps[i]
                bias_correction1 = 1 - beta1**step
                bias_correction2 = 1 - beta2**step

                weight_decay = group["weight_decay"]
                if weight_decay != 0:
                    grad = grad.add(param, alpha=weight_decay)

                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).add_(torch.mul(grad, grad).mul(1 - beta2))

                if group["amsgrad"]:
                    torch.maximum(
                        max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i]
                    )
                    denom = ops.div(
                        ops.sqrt(max_exp_avg_sqs[i]), math.sqrt(bias_correction2)
                    ).add_(group["eps"])
                else:
                    denom = ops.div(
                        ops.sqrt(exp_avg_sq), math.sqrt(bias_correction2)
                    ).add_(group["eps"])
                step_size = group["lr"] / bias_correction1
                param.add_(torch.div(exp_avg, denom).mul(-step_size))

        return loss