Performs a single optimization step. Arguments: - closure (:obj:`callable`, optional): A closure that reevaluates the model and returns the loss.
(self, closure=None)
| 48 | |
| 49 | @torch.no_grad() |
| 50 | def step(self, closure=None): |
| 51 | """Performs a single optimization step. |
| 52 | |
| 53 | Arguments: |
| 54 | - closure (:obj:`callable`, optional): A closure that reevaluates the model and returns the loss. |
| 55 | """ |
| 56 | loss = None |
| 57 | if closure is not None: |
| 58 | loss = closure() |
| 59 | |
| 60 | for group in self.param_groups: |
| 61 | weight_decay = group['weight_decay'] |
| 62 | momentum = group['momentum'] |
| 63 | dampening = group['dampening'] |
| 64 | nesterov = group['nesterov'] |
| 65 | eta = group['eta'] |
| 66 | |
| 67 | for p in group['params']: |
| 68 | if p.grad is None: |
| 69 | continue |
| 70 | d_p = p.grad.data |
| 71 | |
| 72 | # compute local learning rate |
| 73 | weight_norm = torch.norm(p.data) |
| 74 | grad_norm = torch.norm(d_p) |
| 75 | |
| 76 | if weight_decay != 0: |
| 77 | d_p.add_(weight_decay, p.data) |
| 78 | grad_norm.add_(weight_decay, weight_norm) |
| 79 | local_lr = eta * weight_norm / grad_norm |
| 80 | |
| 81 | if momentum != 0: |
| 82 | param_state = self.state[p] |
| 83 | if 'momentum_buffer' not in param_state: |
| 84 | buf = param_state['momentum_buffer'] = torch.zeros_like( |
| 85 | p.data) |
| 86 | buf.mul_(momentum).add_(d_p) |
| 87 | else: |
| 88 | buf = param_state['momentum_buffer'] |
| 89 | buf.mul_(momentum).add_(1 - dampening, d_p) |
| 90 | if nesterov: |
| 91 | d_p = d_p.add(momentum, buf) |
| 92 | else: |
| 93 | d_p = buf |
| 94 | |
| 95 | p.data.add_(-group['lr']*local_lr, d_p) |
| 96 | |
| 97 | return loss |