Update parameters and learning rate.
(self, scaler=None)
| 117 | self.optimizer.zero_grad(set_to_none=set_to_none) |
| 118 | |
| 119 | def step(self, scaler=None): |
| 120 | """Update parameters and learning rate.""" |
| 121 | self._step += 1 |
| 122 | if self.use_dynamic_rate: |
| 123 | rate = self.rate() |
| 124 | for p in self.optimizer.param_groups: |
| 125 | p["lr"] = rate |
| 126 | self._rate = rate |
| 127 | if scaler is None: |
| 128 | self.optimizer.step() |
| 129 | else: |
| 130 | scaler.step(self.optimizer) |
| 131 | else: |
| 132 | if scaler is None: |
| 133 | self.optimizer.step() |
| 134 | else: |
| 135 | scaler.step(self.optimizer) |
| 136 | rates = [] |
| 137 | for p in self.optimizer.param_groups: |
| 138 | rates.append(p["lr"]) |
| 139 | self._rate = sum(rates) / len(rates) |
| 140 | |
| 141 | def rate(self, step: int = None) -> float: |
| 142 | """ |