Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model and returns the loss.
(self, closure=None)
| 89 | |
| 90 | @torch.no_grad() |
| 91 | def step(self, closure=None): |
| 92 | """Performs a single optimization step. |
| 93 | |
| 94 | Arguments: |
| 95 | closure (callable, optional): A closure that reevaluates the model |
| 96 | and returns the loss. |
| 97 | """ |
| 98 | loss = None |
| 99 | if closure is not None: |
| 100 | with torch.enable_grad(): |
| 101 | loss = closure() |
| 102 | |
| 103 | for group in self.param_groups: |
| 104 | weight_decay = group['weight_decay'] |
| 105 | momentum = group['momentum'] |
| 106 | dampening = group['dampening'] |
| 107 | nesterov = group['nesterov'] |
| 108 | |
| 109 | for p in group['params']: |
| 110 | if p.grad is None: |
| 111 | continue |
| 112 | d_p = p.grad |
| 113 | if weight_decay != 0: |
| 114 | d_p = d_p.add(p, alpha=weight_decay) |
| 115 | d_p.mul_(group['lr']) |
| 116 | if momentum != 0: |
| 117 | param_state = self.state[p] |
| 118 | if 'momentum_buffer' not in param_state: |
| 119 | buf = param_state['momentum_buffer'] = torch.clone(d_p).detach() |
| 120 | else: |
| 121 | buf = param_state['momentum_buffer'] |
| 122 | buf.mul_(momentum).add_(d_p, alpha=1 - dampening) |
| 123 | if nesterov: |
| 124 | d_p = d_p.add(buf, alpha=momentum) |
| 125 | else: |
| 126 | d_p = buf |
| 127 | |
| 128 | p.add_(d_p, alpha=-1) |
| 129 | |
| 130 | return loss |
| 131 | |
| 132 | |
| 133 | class TBLog: |