| 14 | |
| 15 | |
| 16 | class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler): |
| 17 | def __init__( |
| 18 | self, |
| 19 | optimizer: torch.optim.Optimizer, |
| 20 | milestones: List[int], |
| 21 | gamma: float = 0.1, |
| 22 | warmup_factor: float = 0.001, |
| 23 | warmup_iters: int = 1000, |
| 24 | warmup_method: str = "linear", |
| 25 | last_epoch: int = -1, |
| 26 | ): |
| 27 | if not list(milestones) == sorted(milestones): |
| 28 | raise ValueError( |
| 29 | "Milestones should be a list of" " increasing integers. Got {}", milestones |
| 30 | ) |
| 31 | self.milestones = milestones |
| 32 | self.gamma = gamma |
| 33 | self.warmup_factor = warmup_factor |
| 34 | self.warmup_iters = warmup_iters |
| 35 | self.warmup_method = warmup_method |
| 36 | super().__init__(optimizer, last_epoch) |
| 37 | |
| 38 | def get_lr(self) -> List[float]: |
| 39 | warmup_factor = _get_warmup_factor_at_iter( |
| 40 | self.warmup_method, self.last_epoch, self.warmup_iters, self.warmup_factor |
| 41 | ) |
| 42 | return [ |
| 43 | base_lr * warmup_factor * self.gamma ** bisect_right(self.milestones, self.last_epoch) |
| 44 | for base_lr in self.base_lrs |
| 45 | ] |
| 46 | |
| 47 | def _compute_values(self) -> List[float]: |
| 48 | # The new interface |
| 49 | return self.get_lr() |
| 50 | |
| 51 | |
| 52 | class WarmupCosineLR(torch.optim.lr_scheduler._LRScheduler): |
no outgoing calls
no test coverage detected