| 161 | return idx |
| 162 | |
| 163 | class CosineWarmupScheduler(LRScheduler): |
| 164 | def __init__(self, optimizer, warmup_iters: int, max_iters: int, initial_lr: float = 1e-10, last_iter: int = -1, min_lr: float = 0.0, decay: bool = True): |
| 165 | self.warmup_iters = warmup_iters |
| 166 | self.max_iters = max_iters |
| 167 | self.initial_lr = initial_lr |
| 168 | self.min_lr = min_lr |
| 169 | self.decay = decay |
| 170 | self._epoch = 0 |
| 171 | super().__init__(optimizer, last_iter) |
| 172 | |
| 173 | def get_lr(self): |
| 174 | # logger.debug(f"step count: {self._step_count} | warmup iters: {self.warmup_iters} | max iters: {self.max_iters}") |
| 175 | if self._step_count <= self.warmup_iters: |
| 176 | return [ |
| 177 | self.initial_lr + (base_lr - self.initial_lr) * self._step_count / self.warmup_iters |
| 178 | for base_lr in self.base_lrs] |
| 179 | elif not self.decay: |
| 180 | return [base_lr for base_lr in self.base_lrs] |
| 181 | elif self._step_count >= self.max_iters: |
| 182 | return [min(self.min_lr, base_lr) for base_lr in self.base_lrs] |
| 183 | else: |
| 184 | cos_iter = self._step_count - self.warmup_iters |
| 185 | cos_max_iter = self.max_iters - self.warmup_iters |
| 186 | cos_theta = cos_iter / cos_max_iter * math.pi |
| 187 | cos_lr = [base_lr * (1. + math.cos(cos_theta)) / 2 for base_lr in self.base_lrs] |
| 188 | return [max(self.min_lr, lr) for lr in cos_lr] |
| 189 | |
| 190 | def step(self, epoch=None): |
| 191 | super().step() |
| 192 | self._epoch = epoch |
| 193 | |
| 194 | |
| 195 | class CosineWeightDecayScheduler(LRScheduler): |