| 217 | super().__init__(optimizer, last_epoch) |
| 218 | |
| 219 | def get_lr(self): |
| 220 | step = self._step_count |
| 221 | # 限制在 [0, max_iters] |
| 222 | if step <= 0: |
| 223 | factor = 0.0 |
| 224 | elif step >= self.max_iters: |
| 225 | factor = 1.0 |
| 226 | else: |
| 227 | # factor 从 0 -> 1，按 1 - cos(pi * t / T) / 2 |
| 228 | factor = 0.5 * (1 - math.cos(math.pi * step / self.max_iters)) |
| 229 | |
| 230 | wd = self.initial_wd + factor * (self.final_wd - self.initial_wd) |
| 231 | return [wd for _ in self.optimizer.param_groups] |
| 232 | |
| 233 | def step(self, epoch=None): |
| 234 | # 先让父类更新 last_epoch |