Spaces:
Running
Running
File size: 827 Bytes
7b0a1ef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
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
|