(self, param_group)
| 51 | self._add_state(param, "step", initializer=0.0) |
| 52 | |
| 53 | def _updates(self, param_group): |
| 54 | lr = param_group["lr"] |
| 55 | lr_decay = param_group["lr_decay"] |
| 56 | weight_decay = param_group["weight_decay"] |
| 57 | eps = param_group["eps"] |
| 58 | |
| 59 | def make_scalar(val): |
| 60 | return tensor(val, dtype="float32") |
| 61 | |
| 62 | # since `conver_inputs` is disabled for param updates, |
| 63 | # scalar should be explicitly tansforred to tensor |
| 64 | |
| 65 | _lr, _lr_decay = map(make_scalar, (lr, lr_decay)) |
| 66 | _weight_decay = make_scalar(weight_decay) |
| 67 | _eps = make_scalar(eps) |
| 68 | |
| 69 | c1, c2, c05 = map(make_scalar, (1.0, 2.0, 0.5)) |
| 70 | |
| 71 | for param in param_group["params"]: |
| 72 | |
| 73 | if param.grad is None: |
| 74 | continue |
| 75 | |
| 76 | states = self._state[param] |
| 77 | step = states["step"] |
| 78 | step += c1 |
| 79 | grad = param.grad |
| 80 | if is_tracing() or weight_decay != 0.0: |
| 81 | grad = grad + param * _weight_decay |
| 82 | |
| 83 | square_avg = states["square_avg"] |
| 84 | square_avg += grad ** c2 |
| 85 | delta = grad / (square_avg + _eps) ** c05 |
| 86 | clr = _lr / (c1 + (step - c1) * _lr_decay) |
| 87 | |
| 88 | param -= clr * delta |
nothing calls this directly
no test coverage detected