Performs a single optimization step. Arguments: closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss.
(self, closure: Callable = None)
| 608 | |
| 609 | @torch.no_grad() |
| 610 | def step(self, closure: Callable = None): |
| 611 | """ |
| 612 | Performs a single optimization step. |
| 613 | |
| 614 | Arguments: |
| 615 | closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss. |
| 616 | """ |
| 617 | loss = None |
| 618 | if closure is not None: |
| 619 | loss = closure() |
| 620 | |
| 621 | for group in self.param_groups: |
| 622 | for p in group["params"]: |
| 623 | if p.grad is None: |
| 624 | continue |
| 625 | grad = p.grad |
| 626 | if grad.is_sparse: |
| 627 | raise RuntimeError("Adam does not support sparse gradients, please consider SparseAdam instead") |
| 628 | |
| 629 | state = self.state[p] |
| 630 | |
| 631 | # State initialization |
| 632 | if len(state) == 0: |
| 633 | state["step"] = 0 |
| 634 | # Exponential moving average of gradient values |
| 635 | state["exp_avg"] = torch.zeros_like(p) |
| 636 | # Exponential moving average of squared gradient values |
| 637 | state["exp_avg_sq"] = torch.zeros_like(p) |
| 638 | |
| 639 | exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"] |
| 640 | beta1, beta2 = group["betas"] |
| 641 | |
| 642 | state["step"] += 1 |
| 643 | |
| 644 | # Decay the first and second moment running average coefficient |
| 645 | # In-place operations to update the averages at the same time |
| 646 | exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1)) |
| 647 | exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) |
| 648 | denom = exp_avg_sq.sqrt().add_(group["eps"]) |
| 649 | |
| 650 | step_size = group["lr"] |
| 651 | if group["correct_bias"]: # No bias correction for Bert |
| 652 | bias_correction1 = 1.0 - beta1 ** state["step"] |
| 653 | bias_correction2 = 1.0 - beta2 ** state["step"] |
| 654 | step_size = step_size * math.sqrt(bias_correction2) / bias_correction1 |
| 655 | |
| 656 | p.addcdiv_(exp_avg, denom, value=-step_size) |
| 657 | |
| 658 | # Just adding the square of the weights to the loss function is *not* |
| 659 | # the correct way of using L2 regularization/weight decay with Adam, |
| 660 | # since that will interact with the m and v parameters in strange ways. |
| 661 | # |
| 662 | # Instead we want to decay the weights in a manner that doesn't interact |
| 663 | # with the m/v parameters. This is equivalent to adding the square |
| 664 | # of the weights to the loss with plain (non-momentum) SGD. |
| 665 | # Add weight decay at the end (fixed version) |
| 666 | if group["weight_decay"] > 0.0: |
| 667 | p.add_(p, alpha=(-group["lr"] * group["weight_decay"])) |
no outgoing calls