Compute the learning rate at a given step. Args: step (int or None): the step to compute the learning rate. If None, use the current step of the optimizer. Returns: learning rate (float)
(self, step: int = None)
| 139 | self._rate = sum(rates) / len(rates) |
| 140 | |
| 141 | def rate(self, step: int = None) -> float: |
| 142 | """ |
| 143 | Compute the learning rate at a given step. |
| 144 | |
| 145 | Args: |
| 146 | step (int or None): |
| 147 | the step to compute the learning rate. |
| 148 | If None, use the current step of the optimizer. |
| 149 | |
| 150 | Returns: |
| 151 | learning rate (float) |
| 152 | """ |
| 153 | if step is None: |
| 154 | step = self._step |
| 155 | return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5))) |
| 156 | |
| 157 | def get_current_lr(self): |
| 158 | """Get the current learning rate.""" |