| 70 | |
| 71 | class StepLRScheduler(_WarmUpLRScheduler): |
| 72 | def __init__(self, optimizer, lr_steps, lr_mults, base_lr, warmup_lr, warmup_steps, last_iter=-1, max_iter=None): |
| 73 | super(StepLRScheduler, self).__init__(optimizer, base_lr, warmup_lr, warmup_steps, last_iter) |
| 74 | |
| 75 | assert len(lr_steps) == len(lr_mults), "{} vs {}".format(milestone, lr_mults) |
| 76 | for x in lr_steps: |
| 77 | assert isinstance(x, int) |
| 78 | if not list(lr_steps) == sorted(lr_steps): |
| 79 | raise ValueError('Milestones should be a list of' |
| 80 | ' increasing integers. Got {}', lr_steps) |
| 81 | self.lr_steps = lr_steps |
| 82 | self.lr_mults = [1.0] |
| 83 | for x in lr_mults: |
| 84 | self.lr_mults.append(self.lr_mults[-1]*x) |
| 85 | |
| 86 | def _get_new_lr(self): |
| 87 | warmup_lr = self._get_warmup_lr() |