| from timm.scheduler.cosine_lr import CosineLRScheduler | |
| from timm.scheduler.step_lr import StepLRScheduler | |
| def build_scheduler(optimizer, n_epoch, n_iter_per_epoch, lr_min=0, warmup_steps=0, warmup_lr_init=0, decay_steps=None, cosine_lr=True): | |
| if decay_steps is None: | |
| decay_steps = n_epoch * n_iter_per_epoch | |
| if cosine_lr: | |
| scheduler = CosineLRScheduler(optimizer, t_initial=decay_steps, lr_min=lr_min, warmup_t=warmup_steps, warmup_lr_init=warmup_lr_init, | |
| cycle_limit=1, t_in_epochs=False, warmup_prefix=True) | |
| else: | |
| scheduler = StepLRScheduler(optimizer, decay_t=decay_steps, warmup_t=warmup_steps, warmup_lr_init=warmup_lr_init, | |
| t_in_epochs=False, warmup_prefix=True) | |
| return scheduler | |