(self, param_group)
| 52 | self._add_state(param, "step", initializer=0.0) |
| 53 | |
| 54 | def _updates(self, param_group): |
| 55 | lr = param_group["lr"] |
| 56 | weight_decay = param_group["weight_decay"] |
| 57 | eps = param_group["eps"] |
| 58 | beta0, beta1 = param_group["betas"] |
| 59 | |
| 60 | def make_scalar(val): |
| 61 | return tensor(val, dtype="float32") |
| 62 | |
| 63 | # since `conver_inputs` is disabled for param updates, |
| 64 | # scalar should be explicitly tansforred to tensor |
| 65 | |
| 66 | _lr, _neg_lr = map(make_scalar, (lr, -lr)) |
| 67 | _weight_decay = make_scalar(weight_decay) |
| 68 | _eps = make_scalar(eps) |
| 69 | _beta0, _beta1 = map(make_scalar, (beta0, beta1)) |
| 70 | |
| 71 | c1, c05 = map(make_scalar, (1.0, 0.5)) |
| 72 | |
| 73 | inplace_mode = int(os.getenv("MEGENGINE_INPLACE_UPDATE", "0")) |
| 74 | if inplace_mode: |
| 75 | # reduce device sync |
| 76 | c1_sub_beta0, c1_sub_beta1 = map(make_scalar, (1 - beta0, 1 - beta1)) |
| 77 | |
| 78 | for param in param_group["params"]: |
| 79 | |
| 80 | if param.grad is None: |
| 81 | continue |
| 82 | |
| 83 | grad = param.grad |
| 84 | if is_tracing() or weight_decay != 0.0: |
| 85 | grad = grad + param * _weight_decay |
| 86 | |
| 87 | states = self._state[param] |
| 88 | |
| 89 | step, exp_avg, exp_avg_sq = ( |
| 90 | states["step"], |
| 91 | states["exp_avg"], |
| 92 | states["exp_avg_sq"], |
| 93 | ) |
| 94 | |
| 95 | if inplace_mode: |
| 96 | _inplace_add_(step, c1, alpha=c1, beta=c1) |
| 97 | _inplace_add_(exp_avg, grad, alpha=_beta0, beta=c1_sub_beta0) |
| 98 | _inplace_add_( |
| 99 | exp_avg_sq, grad * grad, alpha=_beta1, beta=c1_sub_beta1, |
| 100 | ) |
| 101 | |
| 102 | delta = (exp_avg / (c1 - _beta0 ** step)) / ( |
| 103 | (exp_avg_sq / (c1 - _beta1 ** step)) ** c05 + _eps |
| 104 | ) |
| 105 | _inplace_add_(param, delta, alpha=c1, beta=_neg_lr) |
| 106 | continue |
| 107 | |
| 108 | # step = step + c1 |
| 109 | step += c1 |
| 110 | |
| 111 | # exp_avg = _beta0 * exp_avg + grad * (c1 - _beta0) |
nothing calls this directly
no test coverage detected