Spaces:
Running
on
Zero
Running
on
Zero
""" | |
Scheduler | |
Author: Xiaoyang Wu ([email protected]) | |
Please cite our work if the code is helpful to you. | |
""" | |
import torch.optim.lr_scheduler as lr_scheduler | |
from .registry import Registry | |
SCHEDULERS = Registry("schedulers") | |
class MultiStepLR(lr_scheduler.MultiStepLR): | |
def __init__( | |
self, | |
optimizer, | |
milestones, | |
total_steps, | |
gamma=0.1, | |
last_epoch=-1, | |
verbose=False, | |
): | |
super().__init__( | |
optimizer=optimizer, | |
milestones=[rate * total_steps for rate in milestones], | |
gamma=gamma, | |
last_epoch=last_epoch, | |
verbose=verbose, | |
) | |
class MultiStepWithWarmupLR(lr_scheduler.LambdaLR): | |
def __init__( | |
self, | |
optimizer, | |
milestones, | |
total_steps, | |
gamma=0.1, | |
warmup_rate=0.05, | |
warmup_scale=1e-6, | |
last_epoch=-1, | |
verbose=False, | |
): | |
milestones = [rate * total_steps for rate in milestones] | |
def multi_step_with_warmup(s): | |
factor = 1.0 | |
for i in range(len(milestones)): | |
if s < milestones[i]: | |
break | |
factor *= gamma | |
if s <= warmup_rate * total_steps: | |
warmup_coefficient = 1 - (1 - s / warmup_rate / total_steps) * ( | |
1 - warmup_scale | |
) | |
else: | |
warmup_coefficient = 1.0 | |
return warmup_coefficient * factor | |
super().__init__( | |
optimizer=optimizer, | |
lr_lambda=multi_step_with_warmup, | |
last_epoch=last_epoch, | |
verbose=verbose, | |
) | |
class PolyLR(lr_scheduler.LambdaLR): | |
def __init__(self, optimizer, total_steps, power=0.9, last_epoch=-1, verbose=False): | |
super().__init__( | |
optimizer=optimizer, | |
lr_lambda=lambda s: (1 - s / (total_steps + 1)) ** power, | |
last_epoch=last_epoch, | |
verbose=verbose, | |
) | |
class ExpLR(lr_scheduler.LambdaLR): | |
def __init__(self, optimizer, total_steps, gamma=0.9, last_epoch=-1, verbose=False): | |
super().__init__( | |
optimizer=optimizer, | |
lr_lambda=lambda s: gamma ** (s / total_steps), | |
last_epoch=last_epoch, | |
verbose=verbose, | |
) | |
class CosineAnnealingLR(lr_scheduler.CosineAnnealingLR): | |
def __init__(self, optimizer, total_steps, eta_min=0, last_epoch=-1, verbose=False): | |
super().__init__( | |
optimizer=optimizer, | |
T_max=total_steps, | |
eta_min=eta_min, | |
last_epoch=last_epoch, | |
verbose=verbose, | |
) | |
class OneCycleLR(lr_scheduler.OneCycleLR): | |
r""" | |
torch.optim.lr_scheduler.OneCycleLR, Block total_steps | |
""" | |
def __init__( | |
self, | |
optimizer, | |
max_lr, | |
total_steps=None, | |
pct_start=0.3, | |
anneal_strategy="cos", | |
cycle_momentum=True, | |
base_momentum=0.85, | |
max_momentum=0.95, | |
div_factor=25.0, | |
final_div_factor=1e4, | |
three_phase=False, | |
last_epoch=-1, | |
verbose=False, | |
): | |
super().__init__( | |
optimizer=optimizer, | |
max_lr=max_lr, | |
total_steps=total_steps, | |
pct_start=pct_start, | |
anneal_strategy=anneal_strategy, | |
cycle_momentum=cycle_momentum, | |
base_momentum=base_momentum, | |
max_momentum=max_momentum, | |
div_factor=div_factor, | |
final_div_factor=final_div_factor, | |
three_phase=three_phase, | |
last_epoch=last_epoch, | |
verbose=verbose, | |
) | |
def build_scheduler(cfg, optimizer): | |
cfg.optimizer = optimizer | |
return SCHEDULERS.build(cfg=cfg) | |