(self, param_group)
| 60 | self._add_state(param, "momentum_buffer") |
| 61 | |
| 62 | def _updates(self, param_group): |
| 63 | lr = param_group["lr"] |
| 64 | weight_decay = param_group["weight_decay"] |
| 65 | momentum = param_group["momentum"] |
| 66 | |
| 67 | # since `conver_inputs` is disabled for param updates, |
| 68 | # scalar should be explicitly tansforred to tensor |
| 69 | |
| 70 | _lr = tensor(lr, dtype="float32") |
| 71 | _weight_decay = tensor(weight_decay, dtype="float32") |
| 72 | _momentum = tensor(momentum, dtype="float32") |
| 73 | |
| 74 | inplace_mode = int(os.getenv("MEGENGINE_INPLACE_UPDATE", "0")) |
| 75 | if inplace_mode: |
| 76 | _neg_lr = tensor(-lr, dtype="float32") |
| 77 | c1 = tensor(1.0) |
| 78 | |
| 79 | for param in param_group["params"]: |
| 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 | if inplace_mode: |
| 88 | if ( |
| 89 | momentum != tensor(0.0) |
| 90 | if isinstance(momentum, Tensor) |
| 91 | else momentum != 0.0 |
| 92 | ): |
| 93 | v = self._state[param]["momentum_buffer"] |
| 94 | _inplace_add_(v, grad, alpha=_momentum, beta=c1) |
| 95 | if self.nesterov: |
| 96 | grad = grad + v * _momentum |
| 97 | else: |
| 98 | grad = v |
| 99 | _inplace_add_(param, grad, alpha=c1, beta=_neg_lr) |
| 100 | continue |
| 101 | |
| 102 | if ( |
| 103 | momentum != tensor(0.0) |
| 104 | if isinstance(momentum, Tensor) |
| 105 | else momentum != 0.0 |
| 106 | ): |
| 107 | v = self._state[param]["momentum_buffer"] |
| 108 | v *= _momentum |
| 109 | v += grad |
| 110 | if self.nesterov: |
| 111 | grad = grad + v * _momentum |
| 112 | else: |
| 113 | grad = v |
| 114 | param -= _lr * grad |
nothing calls this directly
no test coverage detected