| 20 | |
| 21 | |
| 22 | class WarmUpPolyLR(BaseLR): |
| 23 | def __init__(self, start_lr, lr_power, total_iters, warmup_steps): |
| 24 | self.start_lr = start_lr |
| 25 | self.lr_power = lr_power |
| 26 | self.total_iters = total_iters + 0.0 |
| 27 | self.warmup_steps = warmup_steps |
| 28 | |
| 29 | def get_lr(self, cur_iter): |
| 30 | if cur_iter < self.warmup_steps: |
| 31 | return self.start_lr * (cur_iter / self.warmup_steps) |
| 32 | else: |
| 33 | return self.start_lr * ((1 - float(cur_iter) / self.total_iters) ** self.lr_power) |
| 34 | |
| 35 | |
| 36 | class MultiStageLR(BaseLR): |