mirror of https://github.com/hpcaitech/ColossalAI
78 lines
3.6 KiB
Python
78 lines
3.6 KiB
Python
from torch.optim.lr_scheduler import LambdaLR as _LambdaLR
|
|
from torch.optim.lr_scheduler import MultiplicativeLR as _MultiplicativeLR
|
|
from torch.optim.lr_scheduler import StepLR as _StepLR
|
|
from torch.optim.lr_scheduler import ExponentialLR as _ExponentialLR
|
|
|
|
from colossalai.registry import LR_SCHEDULERS
|
|
|
|
|
|
@LR_SCHEDULERS.register_module
|
|
class LambdaLR(_LambdaLR):
|
|
"""Sets the learning rate of each parameter group to the initial lr
|
|
times a given function. When last_epoch=-1, sets initial lr as lr.
|
|
|
|
Args:
|
|
optimizer (:class:`torch.optim.Optimizer`): Wrapped optimizer.
|
|
total_steps (int): Number of total training steps.
|
|
lr_lambda (Union[``function``, ``list[function]``]): A function which computes a multiplicative
|
|
factor given an integer parameter epoch, or a list of such functions,
|
|
one for each group in optimizer.param_groups, defaults to None.
|
|
last_epoch (int, optional): The index of last epoch, defaults to -1.
|
|
"""
|
|
|
|
def __init__(self, optimizer, total_steps, lr_lambda=None, last_epoch: int = -1) -> None:
|
|
super().__init__(optimizer, lr_lambda, last_epoch=last_epoch)
|
|
|
|
|
|
@LR_SCHEDULERS.register_module
|
|
class MultiplicativeLR(_MultiplicativeLR):
|
|
"""Multiply the learning rate of each parameter group by the factor given
|
|
in the specified function. When last_epoch=-1, sets initial lr as lr.
|
|
|
|
Args:
|
|
optimizer (:class:`torch.optim.Optimizer`): Wrapped optimizer.
|
|
total_steps (int): Number of total training steps.
|
|
lr_lambda (Union[``function``, ``list[function]``]): A function which computes a multiplicative
|
|
factor given an integer parameter epoch, or a list of such functions,
|
|
one for each group in optimizer.param_groups, defaults to None.
|
|
last_epoch (int, optional): The index of last epoch, defaults to -1.
|
|
"""
|
|
|
|
def __init__(self, optimizer, total_steps, lr_lambda=None, last_epoch: int = -1) -> None:
|
|
super().__init__(optimizer, lr_lambda, last_epoch=last_epoch)
|
|
|
|
|
|
@LR_SCHEDULERS.register_module
|
|
class StepLR(_StepLR):
|
|
"""Decays the learning rate of each parameter group by gamma every
|
|
step_size epochs. Notice that such decay can happen simultaneously with
|
|
other changes to the learning rate from outside this scheduler. When
|
|
last_epoch=-1, sets initial lr as lr.
|
|
|
|
Args:
|
|
optimizer (:class:`torch.optim.Optimizer`): Wrapped optimizer.
|
|
total_steps (int): Number of total training steps.
|
|
step_size (int, optional): Period of learning rate decay, defaults to 1.
|
|
gamma (float, optional): Multiplicative factor of learning rate decay, defaults to 0.1.
|
|
last_epoch (int, optional): The index of last epoch, defaults to -1.
|
|
"""
|
|
|
|
def __init__(self, optimizer, total_steps, step_size: int = 1, gamma: float = 0.1, last_epoch: int = -1) -> None:
|
|
super().__init__(optimizer, step_size, gamma=gamma, last_epoch=last_epoch)
|
|
|
|
|
|
@LR_SCHEDULERS.register_module
|
|
class ExponentialLR(_ExponentialLR):
|
|
"""Decays the learning rate of each parameter group by gamma every epoch.
|
|
When last_epoch=-1, sets initial lr as lr
|
|
|
|
Args:
|
|
optimizer (Union[:class:`torch.optim.Optimizer`, :class:`colossalai.nn.optimizer`]): Wrapped optimizer.
|
|
total_steps (int): Number of total training steps.
|
|
gamma (float, optional): Multiplicative factor of learning rate decay, defaults to 1.0.
|
|
last_epoch (int, optional): The index of last epoch, defaults to -1.
|
|
"""
|
|
|
|
def __init__(self, optimizer, total_steps, gamma: float = 1.0, last_epoch: int = -1) -> None:
|
|
super().__init__(optimizer, gamma, last_epoch=last_epoch)
|