INNER CODE UNIT · Python
scheduler_post
QizhiPei/FABind · FABind/fabind/main_fabind.py:273
scheduler_post = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=(args.total_epochs - args.warmup_epochs)*len(train_loader), last_epoch=last_epoch)
elif args.lr_scheduler == "poly_decay":
scheduler_post = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.0, total_iters=(args.total_epochs - args.warmup_epochs)*len(train_loader), last_epoch=last_epoch)
elif args.lr_scheduler == "exp_decay":
scheduler_post = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.995, last_epoch=last_epoch)
elif args.lr_scheduler == "cosine_decay":
scheduler_post = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=(args.total_epochs - args.warmup_epochs)*len(train_loader), eta_min=1e-5, last_epoch=last_epoch)
elif args.lr_scheduler == "cosine_decay_restart":
scheduler_post = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, eta_min=0.0001, last_epoch=last_epoch)
scheduler = torch.optim.lr_scheduler.SequentialLR(
optimizer,
schedulers=[scheduler_warm_up, scheduler_post],
milestones=[args.warmup_epochs*len(train_loader)],
)
(
model,