modules/lr_scheduler.py

from torch.optim.lr_scheduler import *


def make_lr_scheduler(optimizer, lr_scheduler_spec):
    lr_scheduler = {
        'step_lr': StepLR,
        'one_cycle_lr': OneCycleLR,
        'cosine_lr': CosineAnnealingLR,
        'constant_lr': ConstantLR,

    }[lr_scheduler_spec['name']](optimizer, **lr_scheduler_spec['args'])
    return lr_scheduler