Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model and returns the loss.
(self, closure=None)
| 117 | super().__init__(params, defaults) |
| 118 | |
| 119 | def step(self, closure=None): |
| 120 | """Performs a single optimization step. |
| 121 | |
| 122 | Arguments: |
| 123 | closure (callable, optional): A closure that reevaluates the model |
| 124 | and returns the loss. |
| 125 | """ |
| 126 | loss = None |
| 127 | if closure is not None: |
| 128 | loss = closure() |
| 129 | |
| 130 | for group in self.param_groups: |
| 131 | for p in group["params"]: |
| 132 | if p.grad is None: |
| 133 | continue |
| 134 | grad = p.grad.data |
| 135 | if grad.is_sparse: |
| 136 | raise RuntimeError("Adam does not support sparse gradients, please consider SparseAdam instead") |
| 137 | |
| 138 | state = self.state[p] |
| 139 | |
| 140 | # State initialization |
| 141 | if len(state) == 0: |
| 142 | state["step"] = 0 |
| 143 | # Exponential moving average of gradient values |
| 144 | state["exp_avg"] = torch.zeros_like(p.data) |
| 145 | # Exponential moving average of squared gradient values |
| 146 | state["exp_avg_sq"] = torch.zeros_like(p.data) |
| 147 | |
| 148 | exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] |
| 149 | beta1, beta2 = group["betas"] |
| 150 | |
| 151 | state["step"] += 1 |
| 152 | |
| 153 | # Decay the first and second moment running average coefficient |
| 154 | # In-place operations to update the averages at the same time |
| 155 | exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1) |
| 156 | exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) |
| 157 | denom = exp_avg_sq.sqrt().add_(group["eps"]) |
| 158 | |
| 159 | step_size = group["lr"] |
| 160 | if group["correct_bias"]: # No bias correction for Bert |
| 161 | bias_correction1 = 1.0 - beta1 ** state["step"] |
| 162 | bias_correction2 = 1.0 - beta2 ** state["step"] |
| 163 | step_size = step_size * math.sqrt(bias_correction2) / bias_correction1 |
| 164 | |
| 165 | p.data.addcdiv_(exp_avg, denom, value=-step_size) |
| 166 | |
| 167 | # Just adding the square of the weights to the loss function is *not* |
| 168 | # the correct way of using L2 regularization/weight decay with Adam, |
| 169 | # since that will interact with the m and v parameters in strange ways. |
| 170 | # |
| 171 | # Instead we want to decay the weights in a manner that doesn't interact |
| 172 | # with the m/v parameters. This is equivalent to adding the square |
| 173 | # of the weights to the loss with plain (non-momentum) SGD. |
| 174 | # Add weight decay at the end (fixed version) |
| 175 | if group["weight_decay"] > 0.0: |
| 176 | p.data.add_(p.data, alpha=-group["lr"] * group["weight_decay"]) |
no outgoing calls