Spaces:
Runtime error
Runtime error
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. | |
| from bisect import bisect_right | |
| import torch | |
| # FIXME ideally this would be achieved with a CombinedLRScheduler, | |
| # separating MultiStepLR with WarmupLR | |
| # but the current LRScheduler design doesn't allow it | |
| class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler): | |
| def __init__( | |
| self, | |
| optimizer, | |
| milestones, | |
| gamma=0.1, | |
| warmup_factor=1.0 / 3, | |
| warmup_iters=500, | |
| warmup_method="linear", | |
| last_epoch=-1, | |
| pow_schedule_mode = False, | |
| max_iter = 300000, | |
| lr_pow = 0.9 | |
| ): | |
| if not list(milestones) == sorted(milestones): | |
| raise ValueError( | |
| "Milestones should be a list of" " increasing integers. Got {}", | |
| milestones, | |
| ) | |
| if warmup_method not in ("constant", "linear"): | |
| raise ValueError( | |
| "Only 'constant' or 'linear' warmup_method accepted" | |
| "got {}".format(warmup_method) | |
| ) | |
| self.milestones = milestones | |
| self.gamma = gamma | |
| self.warmup_factor = warmup_factor | |
| self.warmup_iters = warmup_iters | |
| self.warmup_method = warmup_method | |
| self.pow_schedule_mode = pow_schedule_mode | |
| self.max_iter = max_iter | |
| self.lr_pow = lr_pow | |
| super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch) | |
| def get_lr(self): | |
| warmup_factor = 1 | |
| if self.last_epoch < self.warmup_iters: | |
| if self.warmup_method == "constant": | |
| warmup_factor = self.warmup_factor | |
| elif self.warmup_method == "linear": | |
| alpha = self.last_epoch / self.warmup_iters | |
| warmup_factor = self.warmup_factor * (1 - alpha) + alpha | |
| if self.pow_schedule_mode: | |
| scale_running_lr = ((1. - float(self.last_epoch) / self.max_iter) ** self.lr_pow) | |
| return [ | |
| base_lr * warmup_factor * scale_running_lr | |
| for base_lr in self.base_lrs | |
| ] | |
| else: | |
| return [ | |
| base_lr | |
| * warmup_factor | |
| * self.gamma ** bisect_right(self.milestones, self.last_epoch) | |
| for base_lr in self.base_lrs | |
| ] | |