Function fetch_optimizer

core/optimizer/__init__.py:4–13

Create the optimizer and learning rate scheduler

(model, cfg)

from torch.optim.lr_scheduler import MultiStepLR, CosineAnnealingLR, ExponentialLR, OneCycleLR

def fetch_optimizer(model, cfg):
    """ Create the optimizer and learning rate scheduler """
    # optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wdecay, eps=args.epsilon)

    # scheduler = optim.lr_scheduler.OneCycleLR(optimizer, args.lr, args.num_steps+100,
    #     pct_start=0.05, cycle_momentum=False, anneal_strategy='linear')
    optimizer = build_optimizer(model, cfg)
    scheduler = build_scheduler(cfg, optimizer)

    return optimizer, scheduler

def build_optimizer(model, config):
    name = config.optimizer