|
|
|
|
|
import math
|
|
|
from typing import Optional
|
|
|
import torch.optim as optim
|
|
|
import torch.optim.lr_scheduler as lr_scheduler
|
|
|
|
|
|
def get_scheduler(
|
|
|
scheduler_name: Optional[str], optimizer: optim.Optimizer, **kwargs
|
|
|
):
|
|
|
|
|
|
def get_warmup_lambda(warm_up_steps, training_steps):
|
|
|
def lr_lambda(steps):
|
|
|
if steps < warm_up_steps:
|
|
|
return (steps + 1) / warm_up_steps
|
|
|
else:
|
|
|
return (training_steps - steps) / (
|
|
|
training_steps - warm_up_steps
|
|
|
)
|
|
|
|
|
|
return lr_lambda
|
|
|
|
|
|
|
|
|
def get_warmup_cosine_lambda(warm_up_steps, training_steps, lr_end):
|
|
|
def lr_lambda(steps):
|
|
|
if steps < warm_up_steps:
|
|
|
return (steps + 1) / warm_up_steps
|
|
|
else:
|
|
|
progress = (steps - warm_up_steps) / (
|
|
|
training_steps - warm_up_steps
|
|
|
)
|
|
|
return lr_end + 0.5 * (1 - lr_end) * (
|
|
|
1 + math.cos(math.pi * progress)
|
|
|
)
|
|
|
|
|
|
return lr_lambda
|
|
|
|
|
|
if scheduler_name is None or scheduler_name.lower() == "constant":
|
|
|
return lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda steps: 1.0)
|
|
|
elif scheduler_name.lower() == "constantwithwarmup":
|
|
|
warm_up_steps = kwargs.get("warm_up_steps", 0)
|
|
|
return lr_scheduler.LambdaLR(
|
|
|
optimizer,
|
|
|
lr_lambda=lambda steps: min(1.0, (steps + 1) / warm_up_steps),
|
|
|
)
|
|
|
else:
|
|
|
raise ValueError(f"Unsupported scheduler: {scheduler_name}") |
