lili_code/runner/scheduler.py

21 lines
354 B
Python

import torch
import math
_scheduler_factory = {
'LambdaLR': torch.optim.lr_scheduler.LambdaLR,
}
def build_scheduler(cfg, optimizer):
assert cfg.scheduler.type in _scheduler_factory
cfg_cp = cfg.scheduler.copy()
cfg_cp.pop('type')
scheduler = _scheduler_factory[cfg.scheduler.type](optimizer, **cfg_cp)
return scheduler