Spaces:
Running on Zero
Running on Zero
File size: 3,074 Bytes
3d736fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import torch
from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
from colossalai.nn.optimizer import HybridAdam
from torch.optim.lr_scheduler import _LRScheduler
def create_optimizer(
model: torch.nn.Module,
optimizer_config: dict,
) -> torch.optim.Optimizer:
"""
Create an optimizer.
Args:
model (torch.nn.Module): The model to be optimized.
optimizer_config (dict): The configuration of the optimizer.
Returns:
torch.optim.Optimizer: The optimizer.
"""
optimizer_name = optimizer_config.pop("cls", "HybridAdam")
if optimizer_name == "HybridAdam":
optimizer_cls = HybridAdam
else:
raise ValueError(f"Unknown optimizer: {optimizer_name}")
optimizer = optimizer_cls(
filter(lambda p: p.requires_grad, model.parameters()),
**optimizer_config,
)
return optimizer
def create_lr_scheduler(
optimizer: torch.optim.Optimizer,
num_steps_per_epoch: int,
epochs: int = 1000,
warmup_steps: int | None = None,
use_cosine_scheduler: bool = False,
initial_lr: float = 1e-6,
) -> _LRScheduler | None:
"""
Create a learning rate scheduler.
Args:
optimizer (torch.optim.Optimizer): The optimizer to be used.
num_steps_per_epoch (int): The number of steps per epoch.
epochs (int): The number of epochs.
warmup_steps (int | None): The number of warmup steps.
use_cosine_scheduler (bool): Whether to use cosine scheduler.
Returns:
_LRScheduler | None: The learning rate scheduler
"""
if warmup_steps is None and not use_cosine_scheduler:
lr_scheduler = None
elif use_cosine_scheduler:
lr_scheduler = CosineAnnealingWarmupLR(
optimizer,
total_steps=num_steps_per_epoch * epochs,
warmup_steps=warmup_steps,
)
else:
lr_scheduler = LinearWarmupLR(optimizer, initial_lr=1e-6, warmup_steps=warmup_steps)
# lr_scheduler = LinearWarmupLR(optimizer, warmup_steps=warmup_steps)
return lr_scheduler
class LinearWarmupLR(_LRScheduler):
"""Linearly warmup learning rate and then linearly decay.
Args:
optimizer (:class:`torch.optim.Optimizer`): Wrapped optimizer.
warmup_steps (int, optional): Number of warmup steps, defaults to 0
last_step (int, optional): The index of last step, defaults to -1. When last_step=-1,
the schedule is started from the beginning or When last_step=-1, sets initial lr as lr.
"""
def __init__(self, optimizer, initial_lr=0, warmup_steps: int = 0, last_epoch: int = -1):
self.initial_lr = initial_lr
self.warmup_steps = warmup_steps
super().__init__(optimizer, last_epoch=last_epoch)
def get_lr(self):
if self.last_epoch < self.warmup_steps:
return [
self.initial_lr + (self.last_epoch + 1) / (self.warmup_steps + 1) * (lr - self.initial_lr)
for lr in self.base_lrs
]
else:
return self.base_lrs
|