2152 lines
83 KiB
Python
2152 lines
83 KiB
Python
# mypy: allow-untyped-defs
|
|
r"""Learning Rate Scheduler."""
|
|
import math
|
|
import types
|
|
import warnings
|
|
from bisect import bisect_right
|
|
from collections import Counter
|
|
from functools import partial, wraps
|
|
from typing import (
|
|
Any,
|
|
Callable,
|
|
cast,
|
|
Dict,
|
|
Iterable,
|
|
List,
|
|
Literal,
|
|
Optional,
|
|
Sequence,
|
|
SupportsFloat,
|
|
TypedDict,
|
|
Union,
|
|
)
|
|
from weakref import ref
|
|
|
|
from torch import inf, Tensor
|
|
|
|
from .optimizer import Optimizer
|
|
|
|
|
|
__all__ = [
|
|
"LambdaLR",
|
|
"MultiplicativeLR",
|
|
"StepLR",
|
|
"MultiStepLR",
|
|
"ConstantLR",
|
|
"LinearLR",
|
|
"ExponentialLR",
|
|
"SequentialLR",
|
|
"CosineAnnealingLR",
|
|
"ChainedScheduler",
|
|
"ReduceLROnPlateau",
|
|
"CyclicLR",
|
|
"CosineAnnealingWarmRestarts",
|
|
"OneCycleLR",
|
|
"PolynomialLR",
|
|
"LRScheduler",
|
|
]
|
|
|
|
EPOCH_DEPRECATION_WARNING = (
|
|
"The epoch parameter in `scheduler.step()` was not necessary and is being "
|
|
"deprecated where possible. Please use `scheduler.step()` to step the "
|
|
"scheduler. During the deprecation, if epoch is different from None, the "
|
|
"closed form is used instead of the new chainable form, where available. "
|
|
"Please open an issue if you are unable to replicate your use case: "
|
|
"https://github.com/pytorch/pytorch/issues/new/choose."
|
|
)
|
|
|
|
|
|
def _check_verbose_deprecated_warning(verbose):
|
|
"""Raise a warning when verbose is not the default value."""
|
|
if verbose != "deprecated":
|
|
warnings.warn(
|
|
"The verbose parameter is deprecated. Please use get_last_lr() "
|
|
"to access the learning rate.",
|
|
UserWarning,
|
|
)
|
|
return verbose
|
|
return False
|
|
|
|
|
|
def _format_param(name: str, optimizer: Optimizer, param):
|
|
"""Return correctly formatted lr/momentum for each param group."""
|
|
|
|
def _copy(_param):
|
|
return _param.clone() if isinstance(_param, Tensor) else _param
|
|
|
|
if isinstance(param, (list, tuple)):
|
|
if len(param) != len(optimizer.param_groups):
|
|
raise ValueError(
|
|
f"{name} must have the same length as optimizer.param_groups. "
|
|
f"{name} has {len(param)} values, param_groups has {len(optimizer.param_groups)}."
|
|
)
|
|
else:
|
|
param = [param] * len(optimizer.param_groups)
|
|
|
|
return list(map(_copy, param))
|
|
|
|
|
|
class LRScheduler:
|
|
r"""Adjusts the learning rate during optimization."""
|
|
|
|
_get_lr_called_within_step: bool = False
|
|
|
|
def __init__(
|
|
self, optimizer: Optimizer, last_epoch=-1, verbose="deprecated"
|
|
): # noqa: D107
|
|
# Attach optimizer
|
|
if not isinstance(optimizer, Optimizer):
|
|
raise TypeError(f"{type(optimizer).__name__} is not an Optimizer")
|
|
self.optimizer = optimizer
|
|
|
|
# Initialize epoch and base learning rates
|
|
if last_epoch == -1:
|
|
for group in optimizer.param_groups:
|
|
initial_lr = group["lr"]
|
|
if isinstance(initial_lr, Tensor):
|
|
initial_lr = initial_lr.clone()
|
|
group.setdefault("initial_lr", initial_lr)
|
|
else:
|
|
for i, group in enumerate(optimizer.param_groups):
|
|
if "initial_lr" not in group:
|
|
raise KeyError(
|
|
"param 'initial_lr' is not specified "
|
|
f"in param_groups[{i}] when resuming an optimizer"
|
|
)
|
|
self.base_lrs: List[float] = [
|
|
group["initial_lr"] for group in optimizer.param_groups
|
|
]
|
|
self.last_epoch = last_epoch
|
|
|
|
# Following https://github.com/pytorch/pytorch/issues/20124
|
|
# We would like to ensure that `lr_scheduler.step()` is called after
|
|
# `optimizer.step()`
|
|
def patch_track_step_called(opt: Optimizer):
|
|
if hasattr(opt.step, "_wrapped_by_lr_sched"):
|
|
# we've already patched
|
|
return opt.step
|
|
|
|
def wrap_step(step_fn):
|
|
opt_ref = ref(self.optimizer)
|
|
func = step_fn.__func__
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
opt = opt_ref()
|
|
opt._opt_called = True # type: ignore[union-attr]
|
|
return func.__get__(opt, opt.__class__)(*args, **kwargs)
|
|
|
|
wrapper._wrapped_by_lr_sched = True # type: ignore[attr-defined]
|
|
return wrapper
|
|
|
|
opt.step = wrap_step(opt.step) # type: ignore[method-assign]
|
|
|
|
patch_track_step_called(self.optimizer)
|
|
self.verbose = _check_verbose_deprecated_warning(verbose)
|
|
self._initial_step()
|
|
|
|
def _initial_step(self):
|
|
"""Initialize step counts and perform a step."""
|
|
self._step_count = 0
|
|
self.step()
|
|
|
|
def state_dict(self):
|
|
"""Return the state of the scheduler as a :class:`dict`.
|
|
|
|
It contains an entry for every variable in self.__dict__ which
|
|
is not the optimizer.
|
|
"""
|
|
return {
|
|
key: value for key, value in self.__dict__.items() if key != "optimizer"
|
|
}
|
|
|
|
def load_state_dict(self, state_dict: Dict[str, Any]):
|
|
"""Load the scheduler's state.
|
|
|
|
Args:
|
|
state_dict (dict): scheduler state. Should be an object returned
|
|
from a call to :meth:`state_dict`.
|
|
"""
|
|
self.__dict__.update(state_dict)
|
|
|
|
def get_last_lr(self) -> List[float]:
|
|
"""Return last computed learning rate by current scheduler."""
|
|
return self._last_lr
|
|
|
|
def get_lr(self) -> List[float]:
|
|
"""Compute learning rate using chainable form of the scheduler."""
|
|
raise NotImplementedError
|
|
|
|
def print_lr(
|
|
self,
|
|
is_verbose: bool,
|
|
group: Dict[str, Any],
|
|
lr: float,
|
|
epoch: Optional[int] = None,
|
|
):
|
|
"""Display the current learning rate.
|
|
|
|
.. deprecated:: 2.4
|
|
``print_lr()`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
"""
|
|
warnings.warn(
|
|
"`LRScheduler.print_lr()` is being deprecated. To fetch the learning rate, "
|
|
"please use `get_last_lr()` instead. For more details, "
|
|
"see https://github.com/pytorch/pytorch/issues/99270.",
|
|
UserWarning,
|
|
)
|
|
if is_verbose:
|
|
if epoch is None:
|
|
print(f"Adjusting learning rate of group {group} to {lr:.4e}.")
|
|
else:
|
|
epoch_str = ("%.2f" if isinstance(epoch, float) else "%.5d") % epoch
|
|
print(
|
|
f"Epoch {epoch_str}: adjusting learning rate of group {group} to {lr:.4e}."
|
|
)
|
|
|
|
def step(self, epoch: Optional[int] = None):
|
|
"""Perform a step."""
|
|
# Raise a warning if old pattern is detected
|
|
# https://github.com/pytorch/pytorch/issues/20124
|
|
if self._step_count == 1:
|
|
if not hasattr(self.optimizer.step, "_wrapped_by_lr_sched"):
|
|
warnings.warn(
|
|
"Seems like `optimizer.step()` has been overridden after learning rate scheduler "
|
|
"initialization. Please, make sure to call `optimizer.step()` before "
|
|
"`lr_scheduler.step()`. See more details at "
|
|
"https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate",
|
|
UserWarning,
|
|
)
|
|
|
|
# Just check if there were two first lr_scheduler.step() calls before optimizer.step()
|
|
elif not getattr(self.optimizer, "_opt_called", False):
|
|
warnings.warn(
|
|
"Detected call of `lr_scheduler.step()` before `optimizer.step()`. "
|
|
"In PyTorch 1.1.0 and later, you should call them in the opposite order: "
|
|
"`optimizer.step()` before `lr_scheduler.step()`. Failure to do this "
|
|
"will result in PyTorch skipping the first value of the learning rate schedule. "
|
|
"See more details at "
|
|
"https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate",
|
|
UserWarning,
|
|
)
|
|
self._step_count += 1
|
|
|
|
with _enable_get_lr_call(self):
|
|
if epoch is None:
|
|
self.last_epoch += 1
|
|
values = self.get_lr()
|
|
else:
|
|
warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)
|
|
self.last_epoch = epoch
|
|
if hasattr(self, "_get_closed_form_lr"):
|
|
values = cast(List[float], self._get_closed_form_lr())
|
|
else:
|
|
values = self.get_lr()
|
|
|
|
for i, data in enumerate(zip(self.optimizer.param_groups, values)):
|
|
param_group, lr = data
|
|
if isinstance(param_group["lr"], Tensor):
|
|
param_group["lr"].fill_(lr)
|
|
else:
|
|
param_group["lr"] = lr
|
|
|
|
self._last_lr: List[float] = [
|
|
group["lr"] for group in self.optimizer.param_groups
|
|
]
|
|
|
|
|
|
def _warn_get_lr_called_within_step(lr_scheduler: LRScheduler):
|
|
if not lr_scheduler._get_lr_called_within_step:
|
|
warnings.warn(
|
|
"To get the last learning rate computed by the scheduler, "
|
|
"please use `get_last_lr()`.",
|
|
UserWarning,
|
|
stacklevel=2,
|
|
)
|
|
|
|
|
|
# Including _LRScheduler for backwards compatibility
|
|
# Subclass instead of assign because we want __name__ of _LRScheduler to be _LRScheduler (assigning would make it LRScheduler).
|
|
class _LRScheduler(LRScheduler):
|
|
pass
|
|
|
|
|
|
class _enable_get_lr_call:
|
|
def __init__(self, o: LRScheduler):
|
|
self.o = o
|
|
|
|
def __enter__(self):
|
|
self.o._get_lr_called_within_step = True
|
|
return self
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
self.o._get_lr_called_within_step = False
|
|
|
|
|
|
class LambdaLR(LRScheduler):
|
|
"""Sets the initial learning rate.
|
|
|
|
The learning rate of each parameter group is set to the initial lr
|
|
times a given function. When last_epoch=-1, sets initial lr as lr.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
lr_lambda (function or list): 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.
|
|
last_epoch (int): The index of last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> # Assuming optimizer has two groups.
|
|
>>> lambda1 = lambda epoch: epoch // 30
|
|
>>> lambda2 = lambda epoch: 0.95 ** epoch
|
|
>>> scheduler = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2])
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
lr_lambda: Union[Callable[[int], float], List[Callable[[int], float]]],
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
self.optimizer = optimizer
|
|
|
|
self.lr_lambdas: List[Callable[[int], float]]
|
|
if not isinstance(lr_lambda, list) and not isinstance(lr_lambda, tuple):
|
|
self.lr_lambdas = [lr_lambda] * len(optimizer.param_groups)
|
|
else:
|
|
if len(lr_lambda) != len(optimizer.param_groups):
|
|
raise ValueError(
|
|
f"Expected {len(optimizer.param_groups)} lr_lambdas, but got {len(lr_lambda)}"
|
|
)
|
|
self.lr_lambdas = list(lr_lambda)
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def state_dict(self):
|
|
"""Return the state of the scheduler as a :class:`dict`.
|
|
|
|
It contains an entry for every variable in self.__dict__ which
|
|
is not the optimizer.
|
|
The learning rate lambda functions will only be saved if they are callable objects
|
|
and not if they are functions or lambdas.
|
|
|
|
When saving or loading the scheduler, please make sure to also save or load the state of the optimizer.
|
|
"""
|
|
state_dict = {
|
|
key: value
|
|
for key, value in self.__dict__.items()
|
|
if key not in ("optimizer", "lr_lambdas")
|
|
}
|
|
state_dict["lr_lambdas"] = [None] * len(self.lr_lambdas)
|
|
|
|
for idx, fn in enumerate(self.lr_lambdas):
|
|
if not isinstance(fn, types.FunctionType):
|
|
state_dict["lr_lambdas"][idx] = fn.__dict__.copy()
|
|
|
|
return state_dict
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Load the scheduler's state.
|
|
|
|
When saving or loading the scheduler, please make sure to also save or load the state of the optimizer.
|
|
|
|
Args:
|
|
state_dict (dict): scheduler state. Should be an object returned
|
|
from a call to :meth:`state_dict`.
|
|
"""
|
|
lr_lambdas = state_dict.pop("lr_lambdas")
|
|
self.__dict__.update(state_dict)
|
|
# Restore state_dict keys in order to prevent side effects
|
|
# https://github.com/pytorch/pytorch/issues/32756
|
|
state_dict["lr_lambdas"] = lr_lambdas
|
|
|
|
for idx, fn in enumerate(lr_lambdas):
|
|
if fn is not None:
|
|
self.lr_lambdas[idx].__dict__.update(fn)
|
|
|
|
def get_lr(self):
|
|
"""Compute learning rate."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
return [
|
|
base_lr * lmbda(self.last_epoch)
|
|
for lmbda, base_lr in zip(self.lr_lambdas, self.base_lrs)
|
|
]
|
|
|
|
|
|
class MultiplicativeLR(LRScheduler):
|
|
"""Multiply the learning rate of each parameter group by the factor given in the specified function.
|
|
|
|
When last_epoch=-1, set initial lr as lr.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
lr_lambda (function or list): 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.
|
|
last_epoch (int): The index of last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> lmbda = lambda epoch: 0.95
|
|
>>> scheduler = MultiplicativeLR(optimizer, lr_lambda=lmbda)
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
lr_lambda: Union[Callable[[int], float], List[Callable[[int], float]]],
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
self.optimizer = optimizer
|
|
|
|
self.lr_lambdas: List[Callable[[int], float]]
|
|
if not isinstance(lr_lambda, list) and not isinstance(lr_lambda, tuple):
|
|
self.lr_lambdas = [lr_lambda] * len(optimizer.param_groups)
|
|
else:
|
|
if len(lr_lambda) != len(optimizer.param_groups):
|
|
raise ValueError(
|
|
f"Expected {len(optimizer.param_groups)} lr_lambdas, but got {len(lr_lambda)}"
|
|
)
|
|
self.lr_lambdas = list(lr_lambda)
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def state_dict(self):
|
|
"""Return the state of the scheduler as a :class:`dict`.
|
|
|
|
It contains an entry for every variable in self.__dict__ which
|
|
is not the optimizer.
|
|
The learning rate lambda functions will only be saved if they are callable objects
|
|
and not if they are functions or lambdas.
|
|
"""
|
|
state_dict = {
|
|
key: value
|
|
for key, value in self.__dict__.items()
|
|
if key not in ("optimizer", "lr_lambdas")
|
|
}
|
|
state_dict["lr_lambdas"] = [None] * len(self.lr_lambdas)
|
|
|
|
for idx, fn in enumerate(self.lr_lambdas):
|
|
if not isinstance(fn, types.FunctionType):
|
|
state_dict["lr_lambdas"][idx] = fn.__dict__.copy()
|
|
|
|
return state_dict
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Load the scheduler's state.
|
|
|
|
Args:
|
|
state_dict (dict): scheduler state. Should be an object returned
|
|
from a call to :meth:`state_dict`.
|
|
"""
|
|
lr_lambdas = state_dict.pop("lr_lambdas")
|
|
self.__dict__.update(state_dict)
|
|
# Restore state_dict keys in order to prevent side effects
|
|
# https://github.com/pytorch/pytorch/issues/32756
|
|
state_dict["lr_lambdas"] = lr_lambdas
|
|
|
|
for idx, fn in enumerate(lr_lambdas):
|
|
if fn is not None:
|
|
self.lr_lambdas[idx].__dict__.update(fn)
|
|
|
|
def get_lr(self):
|
|
"""Compute the learning rate of each parameter group."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
if self.last_epoch > 0:
|
|
return [
|
|
group["lr"] * lmbda(self.last_epoch)
|
|
for lmbda, group in zip(self.lr_lambdas, self.optimizer.param_groups)
|
|
]
|
|
else:
|
|
return [group["lr"] for group in self.optimizer.param_groups]
|
|
|
|
|
|
class StepLR(LRScheduler):
|
|
"""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 (Optimizer): Wrapped optimizer.
|
|
step_size (int): Period of learning rate decay.
|
|
gamma (float): Multiplicative factor of learning rate decay.
|
|
Default: 0.1.
|
|
last_epoch (int): The index of last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> # Assuming optimizer uses lr = 0.05 for all groups
|
|
>>> # lr = 0.05 if epoch < 30
|
|
>>> # lr = 0.005 if 30 <= epoch < 60
|
|
>>> # lr = 0.0005 if 60 <= epoch < 90
|
|
>>> # ...
|
|
>>> scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
step_size: int,
|
|
gamma=0.1,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
self.step_size = step_size
|
|
self.gamma = gamma
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
"""Compute the learning rate of each parameter group."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
if (self.last_epoch == 0) or (self.last_epoch % self.step_size != 0):
|
|
return [group["lr"] for group in self.optimizer.param_groups]
|
|
return [group["lr"] * self.gamma for group in self.optimizer.param_groups]
|
|
|
|
def _get_closed_form_lr(self):
|
|
return [
|
|
base_lr * self.gamma ** (self.last_epoch // self.step_size)
|
|
for base_lr in self.base_lrs
|
|
]
|
|
|
|
|
|
class MultiStepLR(LRScheduler):
|
|
"""Decays the learning rate of each parameter group by gamma once the number of epoch reaches one of the milestones.
|
|
|
|
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 (Optimizer): Wrapped optimizer.
|
|
milestones (list): List of epoch indices. Must be increasing.
|
|
gamma (float): Multiplicative factor of learning rate decay.
|
|
Default: 0.1.
|
|
last_epoch (int): The index of last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> # Assuming optimizer uses lr = 0.05 for all groups
|
|
>>> # lr = 0.05 if epoch < 30
|
|
>>> # lr = 0.005 if 30 <= epoch < 80
|
|
>>> # lr = 0.0005 if epoch >= 80
|
|
>>> scheduler = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
milestones: Iterable[int],
|
|
gamma=0.1,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
self.milestones = Counter(milestones)
|
|
self.gamma = gamma
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
"""Compute the learning rate of each parameter group."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
if self.last_epoch not in self.milestones:
|
|
return [group["lr"] for group in self.optimizer.param_groups]
|
|
return [
|
|
group["lr"] * self.gamma ** self.milestones[self.last_epoch]
|
|
for group in self.optimizer.param_groups
|
|
]
|
|
|
|
def _get_closed_form_lr(self):
|
|
milestones = sorted(self.milestones.elements())
|
|
return [
|
|
base_lr * self.gamma ** bisect_right(milestones, self.last_epoch)
|
|
for base_lr in self.base_lrs
|
|
]
|
|
|
|
|
|
class ConstantLR(LRScheduler):
|
|
"""Multiply the learning rate of each parameter group by a small constant factor.
|
|
|
|
The multiplication is done until the number of epoch reaches a pre-defined milestone: total_iters.
|
|
Notice that such multiplication of the small constant factor 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 (Optimizer): Wrapped optimizer.
|
|
factor (float): The number we multiply learning rate until the milestone. Default: 1./3.
|
|
total_iters (int): The number of steps that the scheduler multiplies the learning rate by the factor.
|
|
Default: 5.
|
|
last_epoch (int): The index of the last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> # Assuming optimizer uses lr = 0.05 for all groups
|
|
>>> # lr = 0.025 if epoch == 0
|
|
>>> # lr = 0.025 if epoch == 1
|
|
>>> # lr = 0.025 if epoch == 2
|
|
>>> # lr = 0.025 if epoch == 3
|
|
>>> # lr = 0.05 if epoch >= 4
|
|
>>> scheduler = ConstantLR(optimizer, factor=0.5, total_iters=4)
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
factor=1.0 / 3,
|
|
total_iters=5,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
if factor > 1.0 or factor < 0:
|
|
raise ValueError(
|
|
"Constant multiplicative factor expected to be between 0 and 1."
|
|
)
|
|
|
|
self.factor = factor
|
|
self.total_iters = total_iters
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
"""Compute the learning rate of each parameter group."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
if self.last_epoch == 0:
|
|
return [group["lr"] * self.factor for group in self.optimizer.param_groups]
|
|
|
|
if self.last_epoch != self.total_iters:
|
|
return [group["lr"] for group in self.optimizer.param_groups]
|
|
|
|
return [
|
|
group["lr"] * (1.0 / self.factor) for group in self.optimizer.param_groups
|
|
]
|
|
|
|
def _get_closed_form_lr(self):
|
|
return [
|
|
base_lr
|
|
* (self.factor + (self.last_epoch >= self.total_iters) * (1 - self.factor))
|
|
for base_lr in self.base_lrs
|
|
]
|
|
|
|
|
|
class LinearLR(LRScheduler):
|
|
"""Decays the learning rate of each parameter group by linearly changing small multiplicative factor.
|
|
|
|
The multiplication is done until the number of epoch reaches a pre-defined milestone: total_iters.
|
|
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 (Optimizer): Wrapped optimizer.
|
|
start_factor (float): The number we multiply learning rate in the first epoch.
|
|
The multiplication factor changes towards end_factor in the following epochs.
|
|
Default: 1./3.
|
|
end_factor (float): The number we multiply learning rate at the end of linear changing
|
|
process. Default: 1.0.
|
|
total_iters (int): The number of iterations that multiplicative factor reaches to 1.
|
|
Default: 5.
|
|
last_epoch (int): The index of the last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> # Assuming optimizer uses lr = 0.05 for all groups
|
|
>>> # lr = 0.025 if epoch == 0
|
|
>>> # lr = 0.03125 if epoch == 1
|
|
>>> # lr = 0.0375 if epoch == 2
|
|
>>> # lr = 0.04375 if epoch == 3
|
|
>>> # lr = 0.05 if epoch >= 4
|
|
>>> scheduler = LinearLR(optimizer, start_factor=0.5, total_iters=4)
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
start_factor=1.0 / 3,
|
|
end_factor=1.0,
|
|
total_iters=5,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
if start_factor > 1.0 or start_factor <= 0:
|
|
raise ValueError(
|
|
"Starting multiplicative factor expected to be greater than 0 and less or equal to 1."
|
|
)
|
|
|
|
if end_factor > 1.0 or end_factor < 0:
|
|
raise ValueError(
|
|
"Ending multiplicative factor expected to be between 0 and 1."
|
|
)
|
|
|
|
self.start_factor = start_factor
|
|
self.end_factor = end_factor
|
|
self.total_iters = total_iters
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
"""Compute the learning rate."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
if self.last_epoch == 0:
|
|
return [
|
|
group["lr"] * self.start_factor for group in self.optimizer.param_groups
|
|
]
|
|
|
|
if self.last_epoch > self.total_iters:
|
|
return [group["lr"] for group in self.optimizer.param_groups]
|
|
|
|
return [
|
|
group["lr"]
|
|
* (
|
|
1.0
|
|
+ (self.end_factor - self.start_factor)
|
|
/ (
|
|
self.total_iters * self.start_factor
|
|
+ (self.last_epoch - 1) * (self.end_factor - self.start_factor)
|
|
)
|
|
)
|
|
for group in self.optimizer.param_groups
|
|
]
|
|
|
|
def _get_closed_form_lr(self):
|
|
return [
|
|
base_lr
|
|
* (
|
|
self.start_factor
|
|
+ (self.end_factor - self.start_factor)
|
|
* min(self.total_iters, self.last_epoch)
|
|
/ self.total_iters
|
|
)
|
|
for base_lr in self.base_lrs
|
|
]
|
|
|
|
|
|
class ExponentialLR(LRScheduler):
|
|
"""Decays the learning rate of each parameter group by gamma every epoch.
|
|
|
|
When last_epoch=-1, sets initial lr as lr.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
gamma (float): Multiplicative factor of learning rate decay.
|
|
last_epoch (int): The index of last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
"""
|
|
|
|
def __init__(
|
|
self, optimizer: Optimizer, gamma: float, last_epoch=-1, verbose="deprecated"
|
|
): # noqa: D107
|
|
self.gamma = gamma
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
"""Compute the learning rate of each parameter group."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
if self.last_epoch == 0:
|
|
return [group["lr"] for group in self.optimizer.param_groups]
|
|
return [group["lr"] * self.gamma for group in self.optimizer.param_groups]
|
|
|
|
def _get_closed_form_lr(self):
|
|
return [base_lr * self.gamma**self.last_epoch for base_lr in self.base_lrs]
|
|
|
|
|
|
class SequentialLR(LRScheduler):
|
|
"""Contains a list of schedulers expected to be called sequentially during the optimization process.
|
|
|
|
Specifically, the schedulers will be called according to the milestone points, which should provide exact
|
|
intervals by which each scheduler should be called at a given epoch.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
schedulers (list): List of chained schedulers.
|
|
milestones (list): List of integers that reflects milestone points.
|
|
last_epoch (int): The index of last epoch. Default: -1.
|
|
verbose (bool | str): Does nothing.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> # Assuming optimizer uses lr = 1. for all groups
|
|
>>> # lr = 0.1 if epoch == 0
|
|
>>> # lr = 0.1 if epoch == 1
|
|
>>> # lr = 0.9 if epoch == 2
|
|
>>> # lr = 0.81 if epoch == 3
|
|
>>> # lr = 0.729 if epoch == 4
|
|
>>> scheduler1 = ConstantLR(optimizer, factor=0.1, total_iters=2)
|
|
>>> scheduler2 = ExponentialLR(optimizer, gamma=0.9)
|
|
>>> scheduler = SequentialLR(optimizer, schedulers=[scheduler1, scheduler2], milestones=[2])
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
schedulers: List[LRScheduler],
|
|
milestones: List[int],
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
if len(schedulers) < 1:
|
|
raise ValueError(
|
|
f"{self.__class__.__name__} expects at least one scheduler, but got no scheduler."
|
|
)
|
|
|
|
for scheduler_idx, scheduler in enumerate(schedulers):
|
|
if not hasattr(scheduler, "optimizer"):
|
|
raise TypeError(
|
|
f"{self.__class__.__name__} at index {scheduler_idx} should have `optimizer` as its attribute."
|
|
)
|
|
if isinstance(scheduler, ReduceLROnPlateau):
|
|
raise ValueError(
|
|
f"{self.__class__.__name__} does not support `ReduceLROnPlateau` scheduler as it "
|
|
"requires additional kwargs to be specified when calling `step`, "
|
|
f"but got one at index {scheduler_idx} in the given schedulers sequence."
|
|
)
|
|
if optimizer != scheduler.optimizer:
|
|
raise ValueError(
|
|
f"{self.__class__.__name__} expects all schedulers to belong to the same optimizer, but "
|
|
f"got scheduler {scheduler.__class__.__name__} at index {scheduler_idx} has {scheduler.optimizer}, "
|
|
f"which is different from {optimizer.__class__.__name__}."
|
|
)
|
|
|
|
if len(milestones) != len(schedulers) - 1:
|
|
raise ValueError(
|
|
"Sequential Schedulers expects number of schedulers provided to be one more "
|
|
f"than the number of milestone points, but got number of schedulers {len(schedulers)} and the "
|
|
f"number of milestones to be equal to {len(milestones)}"
|
|
)
|
|
_check_verbose_deprecated_warning(verbose)
|
|
self._schedulers = schedulers
|
|
self._milestones = milestones
|
|
self.last_epoch = last_epoch + 1
|
|
self.optimizer = optimizer
|
|
|
|
# Reset learning rates back to initial values
|
|
for group in self.optimizer.param_groups:
|
|
group["lr"] = group["initial_lr"]
|
|
|
|
# "Undo" the step performed by other schedulers
|
|
for scheduler in self._schedulers:
|
|
scheduler.last_epoch -= 1
|
|
|
|
# Perform the initial step for only the first scheduler
|
|
self._schedulers[0]._initial_step()
|
|
|
|
self._last_lr = schedulers[0].get_last_lr()
|
|
|
|
def step(self):
|
|
"""Perform a step."""
|
|
self.last_epoch += 1
|
|
idx = bisect_right(self._milestones, self.last_epoch)
|
|
scheduler = self._schedulers[idx]
|
|
if idx > 0 and self._milestones[idx - 1] == self.last_epoch:
|
|
scheduler.step(0)
|
|
else:
|
|
scheduler.step()
|
|
|
|
self._last_lr = scheduler.get_last_lr()
|
|
|
|
def state_dict(self):
|
|
"""Return the state of the scheduler as a :class:`dict`.
|
|
|
|
It contains an entry for every variable in self.__dict__ which
|
|
is not the optimizer.
|
|
The wrapped scheduler states will also be saved.
|
|
"""
|
|
state_dict = {
|
|
key: value
|
|
for key, value in self.__dict__.items()
|
|
if key not in ("optimizer", "_schedulers")
|
|
}
|
|
state_dict["_schedulers"] = [None] * len(self._schedulers)
|
|
|
|
for idx, s in enumerate(self._schedulers):
|
|
state_dict["_schedulers"][idx] = s.state_dict()
|
|
|
|
return state_dict
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Load the scheduler's state.
|
|
|
|
Args:
|
|
state_dict (dict): scheduler state. Should be an object returned
|
|
from a call to :meth:`state_dict`.
|
|
"""
|
|
_schedulers = state_dict.pop("_schedulers")
|
|
self.__dict__.update(state_dict)
|
|
# Restore state_dict keys in order to prevent side effects
|
|
# https://github.com/pytorch/pytorch/issues/32756
|
|
state_dict["_schedulers"] = _schedulers
|
|
|
|
for idx, s in enumerate(_schedulers):
|
|
self._schedulers[idx].load_state_dict(s)
|
|
|
|
|
|
class PolynomialLR(LRScheduler):
|
|
"""Decays the learning rate of each parameter group using a polynomial function in the given total_iters.
|
|
|
|
When last_epoch=-1, sets initial lr as lr.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
total_iters (int): The number of steps that the scheduler decays the learning rate. Default: 5.
|
|
power (float): The power of the polynomial. Default: 1.0.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP("undefined vars")
|
|
>>> # Assuming optimizer uses lr = 0.001 for all groups
|
|
>>> # lr = 0.001 if epoch == 0
|
|
>>> # lr = 0.00075 if epoch == 1
|
|
>>> # lr = 0.00050 if epoch == 2
|
|
>>> # lr = 0.00025 if epoch == 3
|
|
>>> # lr = 0.0 if epoch >= 4
|
|
>>> scheduler = PolynomialLR(optimizer, total_iters=4, power=1.0)
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
total_iters=5,
|
|
power=1.0,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
self.total_iters = total_iters
|
|
self.power = power
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
"""Compute the learning rate."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
if self.last_epoch == 0 or self.last_epoch > self.total_iters:
|
|
return [group["lr"] for group in self.optimizer.param_groups]
|
|
|
|
decay_factor = (
|
|
(1.0 - self.last_epoch / self.total_iters)
|
|
/ (1.0 - (self.last_epoch - 1) / self.total_iters)
|
|
) ** self.power
|
|
return [group["lr"] * decay_factor for group in self.optimizer.param_groups]
|
|
|
|
def _get_closed_form_lr(self):
|
|
return [
|
|
(
|
|
base_lr
|
|
* (1.0 - min(self.total_iters, self.last_epoch) / self.total_iters)
|
|
** self.power
|
|
)
|
|
for base_lr in self.base_lrs
|
|
]
|
|
|
|
|
|
class CosineAnnealingLR(LRScheduler):
|
|
r"""Set the learning rate of each parameter group using a cosine annealing schedule.
|
|
|
|
The :math:`\eta_{max}` is set to the initial lr and
|
|
:math:`T_{cur}` is the number of epochs since the last restart in SGDR:
|
|
|
|
.. math::
|
|
\begin{aligned}
|
|
\eta_t & = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1
|
|
+ \cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right),
|
|
& T_{cur} \neq (2k+1)T_{max}; \\
|
|
\eta_{t+1} & = \eta_{t} + \frac{1}{2}(\eta_{max} - \eta_{min})
|
|
\left(1 - \cos\left(\frac{1}{T_{max}}\pi\right)\right),
|
|
& T_{cur} = (2k+1)T_{max}.
|
|
\end{aligned}
|
|
|
|
When last_epoch=-1, sets initial lr as lr. Notice that because the schedule
|
|
is defined recursively, the learning rate can be simultaneously modified
|
|
outside this scheduler by other operators. If the learning rate is set
|
|
solely by this scheduler, the learning rate at each step becomes:
|
|
|
|
.. math::
|
|
\eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 +
|
|
\cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right)
|
|
|
|
It has been proposed in
|
|
`SGDR: Stochastic Gradient Descent with Warm Restarts`_. Note that this only
|
|
implements the cosine annealing part of SGDR, and not the restarts.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
T_max (int): Maximum number of iterations.
|
|
eta_min (float): Minimum learning rate. Default: 0.
|
|
last_epoch (int): The index of last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
.. _SGDR\: Stochastic Gradient Descent with Warm Restarts:
|
|
https://arxiv.org/abs/1608.03983
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
T_max: int,
|
|
eta_min=0.0,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
self.T_max = T_max
|
|
self.eta_min = eta_min
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
"""Retrieve the learning rate of each parameter group."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
if self.last_epoch == 0:
|
|
return [group["lr"] for group in self.optimizer.param_groups]
|
|
elif self._step_count == 1 and self.last_epoch > 0:
|
|
return [
|
|
self.eta_min
|
|
+ (base_lr - self.eta_min)
|
|
* (1 + math.cos((self.last_epoch) * math.pi / self.T_max))
|
|
/ 2
|
|
for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
|
|
]
|
|
elif (self.last_epoch - 1 - self.T_max) % (2 * self.T_max) == 0:
|
|
return [
|
|
group["lr"]
|
|
+ (base_lr - self.eta_min) * (1 - math.cos(math.pi / self.T_max)) / 2
|
|
for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
|
|
]
|
|
return [
|
|
(1 + math.cos(math.pi * self.last_epoch / self.T_max))
|
|
/ (1 + math.cos(math.pi * (self.last_epoch - 1) / self.T_max))
|
|
* (group["lr"] - self.eta_min)
|
|
+ self.eta_min
|
|
for group in self.optimizer.param_groups
|
|
]
|
|
|
|
def _get_closed_form_lr(self):
|
|
return [
|
|
self.eta_min
|
|
+ (base_lr - self.eta_min)
|
|
* (1 + math.cos(math.pi * self.last_epoch / self.T_max))
|
|
/ 2
|
|
for base_lr in self.base_lrs
|
|
]
|
|
|
|
|
|
class ChainedScheduler(LRScheduler):
|
|
"""Chains a list of learning rate schedulers.
|
|
|
|
Takes in a sequence of chainable learning rate schedulers and calls their
|
|
step() functions consecutively in just one call to step().
|
|
|
|
Args:
|
|
schedulers (sequence): sequence of chained schedulers.
|
|
optimizer (Optimizer, optional): Wrapped optimizer. Default: None.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> # Assuming optimizer uses lr = 1. for all groups
|
|
>>> # lr = 0.09 if epoch == 0
|
|
>>> # lr = 0.081 if epoch == 1
|
|
>>> # lr = 0.729 if epoch == 2
|
|
>>> # lr = 0.6561 if epoch == 3
|
|
>>> # lr = 0.59049 if epoch >= 4
|
|
>>> scheduler1 = ConstantLR(optimizer, factor=0.1, total_iters=2)
|
|
>>> scheduler2 = ExponentialLR(optimizer, gamma=0.9)
|
|
>>> scheduler = ChainedScheduler([scheduler1, scheduler2], optimizer=optimizer)
|
|
>>> for epoch in range(100):
|
|
>>> train(...)
|
|
>>> validate(...)
|
|
>>> scheduler.step()
|
|
"""
|
|
|
|
def __init__(
|
|
self, schedulers: Sequence[LRScheduler], optimizer: Optional[Optimizer] = None
|
|
): # noqa: D107
|
|
if len(schedulers) < 1:
|
|
raise ValueError(
|
|
f"{self.__class__.__name__} expects at least one scheduler to be chained, but got no scheduler."
|
|
)
|
|
|
|
optimizer = optimizer or schedulers[0].optimizer
|
|
for scheduler_idx, scheduler in enumerate(schedulers):
|
|
if not hasattr(scheduler, "optimizer"):
|
|
raise TypeError(
|
|
f"{self.__class__.__name__} at index {scheduler_idx} should have `optimizer` as its attribute."
|
|
)
|
|
if isinstance(scheduler, ReduceLROnPlateau):
|
|
raise ValueError(
|
|
f"{self.__class__.__name__} does not support `ReduceLROnPlateau` scheduler as it "
|
|
"requires additional kwargs to be specified when calling `step`, "
|
|
f"but got one at index {scheduler_idx} in the given schedulers sequence."
|
|
)
|
|
if optimizer != scheduler.optimizer:
|
|
raise ValueError(
|
|
f"{self.__class__.__name__} expects all schedulers to belong to the same optimizer, but "
|
|
f"got scheduler {scheduler.__class__.__name__} at index {scheduler_idx} has {scheduler.optimizer}, "
|
|
f"which is different from {optimizer.__class__.__name__}."
|
|
)
|
|
self._schedulers = schedulers
|
|
self.optimizer = optimizer
|
|
self._last_lr = [
|
|
group["lr"] for group in self._schedulers[-1].optimizer.param_groups
|
|
]
|
|
|
|
def step(self):
|
|
"""Perform a step."""
|
|
for scheduler in self._schedulers:
|
|
scheduler.step()
|
|
self._last_lr = [
|
|
group["lr"] for group in self._schedulers[-1].optimizer.param_groups
|
|
]
|
|
|
|
def state_dict(self):
|
|
"""Return the state of the scheduler as a :class:`dict`.
|
|
|
|
It contains an entry for every variable in self.__dict__ which
|
|
is not the optimizer.
|
|
The wrapped scheduler states will also be saved.
|
|
"""
|
|
state_dict = {
|
|
key: value
|
|
for key, value in self.__dict__.items()
|
|
if key not in ("optimizer", "_schedulers")
|
|
}
|
|
state_dict["_schedulers"] = [None] * len(self._schedulers)
|
|
|
|
for idx, s in enumerate(self._schedulers):
|
|
state_dict["_schedulers"][idx] = s.state_dict()
|
|
|
|
return state_dict
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Load the scheduler's state.
|
|
|
|
Args:
|
|
state_dict (dict): scheduler state. Should be an object returned
|
|
from a call to :meth:`state_dict`.
|
|
"""
|
|
_schedulers = state_dict.pop("_schedulers")
|
|
self.__dict__.update(state_dict)
|
|
# Restore state_dict keys in order to prevent side effects
|
|
# https://github.com/pytorch/pytorch/issues/32756
|
|
state_dict["_schedulers"] = _schedulers
|
|
|
|
for idx, s in enumerate(_schedulers):
|
|
self._schedulers[idx].load_state_dict(s)
|
|
|
|
|
|
class ReduceLROnPlateau(LRScheduler):
|
|
"""Reduce learning rate when a metric has stopped improving.
|
|
|
|
Models often benefit from reducing the learning rate by a factor
|
|
of 2-10 once learning stagnates. This scheduler reads a metrics
|
|
quantity and if no improvement is seen for a 'patience' number
|
|
of epochs, the learning rate is reduced.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
mode (str): One of `min`, `max`. In `min` mode, lr will
|
|
be reduced when the quantity monitored has stopped
|
|
decreasing; in `max` mode it will be reduced when the
|
|
quantity monitored has stopped increasing. Default: 'min'.
|
|
factor (float): Factor by which the learning rate will be
|
|
reduced. new_lr = lr * factor. Default: 0.1.
|
|
patience (int): The number of allowed epochs with no improvement after
|
|
which the learning rate will be reduced.
|
|
For example, consider the case of having no patience (`patience = 0`).
|
|
In the first epoch, a baseline is established and is always considered good as there's no previous baseline.
|
|
In the second epoch, if the performance is worse than the baseline,
|
|
we have what is considered an intolerable epoch.
|
|
Since the count of intolerable epochs (1) is greater than the patience level (0),
|
|
the learning rate is reduced at the end of this epoch.
|
|
From the third epoch onwards, the learning rate continues to be reduced at the end of each epoch
|
|
if the performance is worse than the baseline. If the performance improves or remains the same,
|
|
the learning rate is not adjusted.
|
|
Default: 10.
|
|
threshold (float): Threshold for measuring the new optimum,
|
|
to only focus on significant changes. Default: 1e-4.
|
|
threshold_mode (str): One of `rel`, `abs`. In `rel` mode,
|
|
dynamic_threshold = best * ( 1 + threshold ) in 'max'
|
|
mode or best * ( 1 - threshold ) in `min` mode.
|
|
In `abs` mode, dynamic_threshold = best + threshold in
|
|
`max` mode or best - threshold in `min` mode. Default: 'rel'.
|
|
cooldown (int): Number of epochs to wait before resuming
|
|
normal operation after lr has been reduced. Default: 0.
|
|
min_lr (float or list): A scalar or a list of scalars. A
|
|
lower bound on the learning rate of all param groups
|
|
or each group respectively. Default: 0.
|
|
eps (float): Minimal decay applied to lr. If the difference
|
|
between new and old lr is smaller than eps, the update is
|
|
ignored. Default: 1e-8.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
|
|
>>> scheduler = ReduceLROnPlateau(optimizer, 'min')
|
|
>>> for epoch in range(10):
|
|
>>> train(...)
|
|
>>> val_loss = validate(...)
|
|
>>> # Note that step should be called after validate()
|
|
>>> scheduler.step(val_loss)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
mode: Literal["min", "max"] = "min",
|
|
factor=0.1,
|
|
patience=10,
|
|
threshold=1e-4,
|
|
threshold_mode: Literal["rel", "abs"] = "rel",
|
|
cooldown=0,
|
|
min_lr: Union[List[float], float] = 0,
|
|
eps=1e-8,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
if factor >= 1.0:
|
|
raise ValueError("Factor should be < 1.0.")
|
|
self.factor = factor
|
|
|
|
# Attach optimizer
|
|
if not isinstance(optimizer, Optimizer):
|
|
raise TypeError(f"{type(optimizer).__name__} is not an Optimizer")
|
|
self.optimizer = optimizer
|
|
|
|
if isinstance(min_lr, (list, tuple)):
|
|
if len(min_lr) != len(optimizer.param_groups):
|
|
raise ValueError(
|
|
f"expected {len(optimizer.param_groups)} min_lrs, got {len(min_lr)}"
|
|
)
|
|
self.min_lrs = list(min_lr)
|
|
else:
|
|
self.min_lrs = [min_lr] * len(optimizer.param_groups)
|
|
|
|
self.patience = patience
|
|
|
|
self.verbose = _check_verbose_deprecated_warning(verbose)
|
|
self.cooldown = cooldown
|
|
self.cooldown_counter = 0
|
|
self.mode = mode
|
|
self.threshold = threshold
|
|
self.threshold_mode = threshold_mode
|
|
self.best: float
|
|
self.num_bad_epochs: int
|
|
self.mode_worse: float # the worse value for the chosen mode
|
|
self.eps = eps
|
|
self.last_epoch = 0
|
|
self._last_lr = [group["lr"] for group in self.optimizer.param_groups]
|
|
self._init_is_better(
|
|
mode=mode, threshold=threshold, threshold_mode=threshold_mode
|
|
)
|
|
self._reset()
|
|
|
|
def _reset(self):
|
|
"""Reset num_bad_epochs counter and cooldown counter."""
|
|
self.best = self.mode_worse
|
|
self.cooldown_counter = 0
|
|
self.num_bad_epochs = 0
|
|
|
|
def step(self, metrics: SupportsFloat, epoch=None): # type: ignore[override]
|
|
"""Perform a step."""
|
|
# convert `metrics` to float, in case it's a zero-dim Tensor
|
|
current = float(metrics)
|
|
if epoch is None:
|
|
epoch = self.last_epoch + 1
|
|
else:
|
|
warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning)
|
|
self.last_epoch = epoch
|
|
|
|
if self.is_better(current, self.best):
|
|
self.best = current
|
|
self.num_bad_epochs = 0
|
|
else:
|
|
self.num_bad_epochs += 1
|
|
|
|
if self.in_cooldown:
|
|
self.cooldown_counter -= 1
|
|
self.num_bad_epochs = 0 # ignore any bad epochs in cooldown
|
|
|
|
if self.num_bad_epochs > self.patience:
|
|
self._reduce_lr(epoch)
|
|
self.cooldown_counter = self.cooldown
|
|
self.num_bad_epochs = 0
|
|
|
|
self._last_lr = [group["lr"] for group in self.optimizer.param_groups]
|
|
|
|
def _reduce_lr(self, epoch):
|
|
for i, param_group in enumerate(self.optimizer.param_groups):
|
|
old_lr = float(param_group["lr"])
|
|
new_lr = max(old_lr * self.factor, self.min_lrs[i])
|
|
if old_lr - new_lr > self.eps:
|
|
param_group["lr"] = new_lr
|
|
|
|
@property
|
|
def in_cooldown(self): # noqa: D102
|
|
return self.cooldown_counter > 0
|
|
|
|
def is_better(self, a, best): # noqa: D102
|
|
if self.mode == "min" and self.threshold_mode == "rel":
|
|
rel_epsilon = 1.0 - self.threshold
|
|
return a < best * rel_epsilon
|
|
|
|
elif self.mode == "min" and self.threshold_mode == "abs":
|
|
return a < best - self.threshold
|
|
|
|
elif self.mode == "max" and self.threshold_mode == "rel":
|
|
rel_epsilon = self.threshold + 1.0
|
|
return a > best * rel_epsilon
|
|
|
|
else: # mode == 'max' and epsilon_mode == 'abs':
|
|
return a > best + self.threshold
|
|
|
|
def _init_is_better(self, mode, threshold, threshold_mode):
|
|
if mode not in {"min", "max"}:
|
|
raise ValueError("mode " + mode + " is unknown!")
|
|
if threshold_mode not in {"rel", "abs"}:
|
|
raise ValueError("threshold mode " + threshold_mode + " is unknown!")
|
|
|
|
if mode == "min":
|
|
self.mode_worse = inf
|
|
else: # mode == 'max':
|
|
self.mode_worse = -inf
|
|
|
|
self.mode = mode
|
|
self.threshold = threshold
|
|
self.threshold_mode = threshold_mode
|
|
|
|
def state_dict(self): # noqa: D102
|
|
return {
|
|
key: value for key, value in self.__dict__.items() if key != "optimizer"
|
|
}
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Load the scheduler's state."""
|
|
self.__dict__.update(state_dict)
|
|
self._init_is_better(
|
|
mode=self.mode, threshold=self.threshold, threshold_mode=self.threshold_mode
|
|
)
|
|
|
|
|
|
class CyclicLR(LRScheduler):
|
|
r"""Sets the learning rate of each parameter group according to cyclical learning rate policy (CLR).
|
|
|
|
The policy cycles the learning rate between two boundaries with a constant frequency,
|
|
as detailed in the paper `Cyclical Learning Rates for Training Neural Networks`_.
|
|
The distance between the two boundaries can be scaled on a per-iteration
|
|
or per-cycle basis.
|
|
|
|
Cyclical learning rate policy changes the learning rate after every batch.
|
|
`step` should be called after a batch has been used for training.
|
|
|
|
This class has three built-in policies, as put forth in the paper:
|
|
|
|
* "triangular": A basic triangular cycle without amplitude scaling.
|
|
* "triangular2": A basic triangular cycle that scales initial amplitude by half each cycle.
|
|
* "exp_range": A cycle that scales initial amplitude by :math:`\text{gamma}^{\text{cycle iterations}}`
|
|
at each cycle iteration.
|
|
|
|
This implementation was adapted from the github repo: `bckenstler/CLR`_
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
base_lr (float or list): Initial learning rate which is the
|
|
lower boundary in the cycle for each parameter group.
|
|
max_lr (float or list): Upper learning rate boundaries in the cycle
|
|
for each parameter group. Functionally,
|
|
it defines the cycle amplitude (max_lr - base_lr).
|
|
The lr at any cycle is the sum of base_lr
|
|
and some scaling of the amplitude; therefore
|
|
max_lr may not actually be reached depending on
|
|
scaling function.
|
|
step_size_up (int): Number of training iterations in the
|
|
increasing half of a cycle. Default: 2000
|
|
step_size_down (int): Number of training iterations in the
|
|
decreasing half of a cycle. If step_size_down is None,
|
|
it is set to step_size_up. Default: None
|
|
mode (str): One of {triangular, triangular2, exp_range}.
|
|
Values correspond to policies detailed above.
|
|
If scale_fn is not None, this argument is ignored.
|
|
Default: 'triangular'
|
|
gamma (float): Constant in 'exp_range' scaling function:
|
|
gamma**(cycle iterations)
|
|
Default: 1.0
|
|
scale_fn (function): Custom scaling policy defined by a single
|
|
argument lambda function, where
|
|
0 <= scale_fn(x) <= 1 for all x >= 0.
|
|
If specified, then 'mode' is ignored.
|
|
Default: None
|
|
scale_mode (str): {'cycle', 'iterations'}.
|
|
Defines whether scale_fn is evaluated on
|
|
cycle number or cycle iterations (training
|
|
iterations since start of cycle).
|
|
Default: 'cycle'
|
|
cycle_momentum (bool): If ``True``, momentum is cycled inversely
|
|
to learning rate between 'base_momentum' and 'max_momentum'.
|
|
Default: True
|
|
base_momentum (float or list): Lower momentum boundaries in the cycle
|
|
for each parameter group. Note that momentum is cycled inversely
|
|
to learning rate; at the peak of a cycle, momentum is
|
|
'base_momentum' and learning rate is 'max_lr'.
|
|
Default: 0.8
|
|
max_momentum (float or list): Upper momentum boundaries in the cycle
|
|
for each parameter group. Functionally,
|
|
it defines the cycle amplitude (max_momentum - base_momentum).
|
|
The momentum at any cycle is the difference of max_momentum
|
|
and some scaling of the amplitude; therefore
|
|
base_momentum may not actually be reached depending on
|
|
scaling function. Note that momentum is cycled inversely
|
|
to learning rate; at the start of a cycle, momentum is 'max_momentum'
|
|
and learning rate is 'base_lr'
|
|
Default: 0.9
|
|
last_epoch (int): The index of the last batch. This parameter is used when
|
|
resuming a training job. Since `step()` should be invoked after each
|
|
batch instead of after each epoch, this number represents the total
|
|
number of *batches* computed, not the total number of epochs computed.
|
|
When last_epoch=-1, the schedule is started from the beginning.
|
|
Default: -1
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
|
|
>>> scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.01, max_lr=0.1)
|
|
>>> data_loader = torch.utils.data.DataLoader(...)
|
|
>>> for epoch in range(10):
|
|
>>> for batch in data_loader:
|
|
>>> train_batch(...)
|
|
>>> scheduler.step()
|
|
|
|
|
|
.. _Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186
|
|
.. _bckenstler/CLR: https://github.com/bckenstler/CLR
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
base_lr: Union[float, List[float]],
|
|
max_lr: Union[float, List[float]],
|
|
step_size_up=2000,
|
|
step_size_down: Optional[int] = None,
|
|
mode: Literal["triangular", "triangular2", "exp_range"] = "triangular",
|
|
gamma=1.0,
|
|
scale_fn: Optional[Callable[[float], float]] = None,
|
|
scale_mode: Literal["cycle", "iterations"] = "cycle",
|
|
cycle_momentum=True,
|
|
base_momentum=0.8,
|
|
max_momentum=0.9,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
# Attach optimizer
|
|
if not isinstance(optimizer, Optimizer):
|
|
raise TypeError(f"{type(optimizer).__name__} is not an Optimizer")
|
|
self.optimizer = optimizer
|
|
|
|
base_lrs = _format_param("base_lr", optimizer, base_lr)
|
|
if last_epoch == -1:
|
|
for lr, group in zip(base_lrs, optimizer.param_groups):
|
|
if isinstance(group["lr"], Tensor):
|
|
lr_val = lr.item() if isinstance(lr, Tensor) else lr
|
|
group["lr"].fill_(lr_val)
|
|
else:
|
|
group["lr"] = lr
|
|
|
|
self.max_lrs = _format_param("max_lr", optimizer, max_lr)
|
|
|
|
step_size_up = float(step_size_up)
|
|
step_size_down = (
|
|
float(step_size_down) if step_size_down is not None else step_size_up
|
|
)
|
|
self.total_size = step_size_up + step_size_down
|
|
self.step_ratio = step_size_up / self.total_size
|
|
|
|
if mode not in ["triangular", "triangular2", "exp_range"] and scale_fn is None:
|
|
raise ValueError("mode is invalid and scale_fn is None")
|
|
|
|
self.mode = mode
|
|
self.gamma = gamma
|
|
|
|
self._scale_fn_ref: Callable[[float], float]
|
|
self._scale_fn_custom = scale_fn
|
|
self.scale_mode = scale_mode
|
|
self._init_scale_fn()
|
|
|
|
self.cycle_momentum = cycle_momentum
|
|
if cycle_momentum:
|
|
if (
|
|
"momentum" not in optimizer.defaults
|
|
and "betas" not in optimizer.defaults
|
|
):
|
|
raise ValueError(
|
|
"optimizer must support momentum or beta1 with `cycle_momentum` option enabled"
|
|
)
|
|
|
|
self.use_beta1 = "betas" in self.optimizer.defaults
|
|
self.base_momentums = _format_param(
|
|
"base_momentum", optimizer, base_momentum
|
|
)
|
|
self.max_momentums = _format_param("max_momentum", optimizer, max_momentum)
|
|
if last_epoch == -1:
|
|
for m_momentum, b_momentum, group in zip(
|
|
self.max_momentums, self.base_momentums, optimizer.param_groups
|
|
):
|
|
if self.use_beta1:
|
|
group["betas"] = (m_momentum, *group["betas"][1:])
|
|
else:
|
|
group["momentum"] = m_momentum
|
|
group["max_momentum"] = m_momentum
|
|
group["base_momentum"] = b_momentum
|
|
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
self.base_lrs = base_lrs
|
|
|
|
def _init_scale_fn(self):
|
|
if self._scale_fn_custom is not None:
|
|
return
|
|
if self.mode == "triangular":
|
|
self._scale_fn_ref = self._triangular_scale_fn
|
|
self.scale_mode = "cycle"
|
|
elif self.mode == "triangular2":
|
|
self._scale_fn_ref = self._triangular2_scale_fn
|
|
self.scale_mode = "cycle"
|
|
elif self.mode == "exp_range":
|
|
self._scale_fn_ref = partial(self._exp_range_scale_fn, self.gamma)
|
|
self.scale_mode = "iterations"
|
|
|
|
def scale_fn(self, x) -> float:
|
|
"""Get the scaling policy."""
|
|
if self._scale_fn_custom is not None:
|
|
return self._scale_fn_custom(x)
|
|
else:
|
|
return self._scale_fn_ref(x) # static method
|
|
|
|
@staticmethod
|
|
def _triangular_scale_fn(x: float) -> float:
|
|
return 1.0
|
|
|
|
@staticmethod
|
|
def _triangular2_scale_fn(x: float) -> float:
|
|
return 1 / (2.0 ** (x - 1))
|
|
|
|
@staticmethod
|
|
def _exp_range_scale_fn(gamma: float, x: float) -> float:
|
|
return gamma**x
|
|
|
|
def get_lr(self):
|
|
"""Calculate the learning rate at batch index.
|
|
|
|
This function treats `self.last_epoch` as the last batch index.
|
|
|
|
If `self.cycle_momentum` is ``True``, this function has a side effect of
|
|
updating the optimizer's momentum.
|
|
"""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
cycle = math.floor(1 + self.last_epoch / self.total_size)
|
|
x = 1.0 + self.last_epoch / self.total_size - cycle
|
|
if x <= self.step_ratio:
|
|
scale_factor = x / self.step_ratio
|
|
else:
|
|
scale_factor = (x - 1) / (self.step_ratio - 1)
|
|
|
|
lrs = []
|
|
for base_lr, max_lr in zip(self.base_lrs, self.max_lrs):
|
|
base_height = (max_lr - base_lr) * scale_factor
|
|
if self.scale_mode == "cycle":
|
|
lr = base_lr + base_height * self.scale_fn(cycle)
|
|
else:
|
|
lr = base_lr + base_height * self.scale_fn(self.last_epoch)
|
|
lrs.append(lr)
|
|
|
|
if self.cycle_momentum:
|
|
momentums = []
|
|
for base_momentum, max_momentum in zip(
|
|
self.base_momentums, self.max_momentums
|
|
):
|
|
base_height = (max_momentum - base_momentum) * scale_factor
|
|
if self.scale_mode == "cycle":
|
|
momentum = max_momentum - base_height * self.scale_fn(cycle)
|
|
else:
|
|
momentum = max_momentum - base_height * self.scale_fn(
|
|
self.last_epoch
|
|
)
|
|
momentums.append(momentum)
|
|
for param_group, momentum in zip(self.optimizer.param_groups, momentums):
|
|
if self.use_beta1:
|
|
param_group["betas"] = (momentum, *param_group["betas"][1:])
|
|
else:
|
|
param_group["momentum"] = momentum
|
|
|
|
return lrs
|
|
|
|
def state_dict(self): # noqa: D102
|
|
state = super().state_dict()
|
|
# We are dropping the `_scale_fn_ref` attribute because it is a
|
|
# `weakref.WeakMethod` and can't be pickled.
|
|
state.pop("_scale_fn_ref", None)
|
|
fn = state.pop("_scale_fn_custom")
|
|
state["_scale_fn_custom"] = None
|
|
if fn is not None and not isinstance(fn, types.FunctionType):
|
|
# The _scale_fn_custom will only be saved if it is a callable object
|
|
# and not if it is a function or lambda.
|
|
state["_scale_fn_custom"] = fn.__dict__.copy()
|
|
|
|
return state
|
|
|
|
def load_state_dict(self, state_dict):
|
|
"""Load the scheduler's state."""
|
|
fn = state_dict.pop("_scale_fn_custom")
|
|
super().load_state_dict(state_dict)
|
|
if fn is not None:
|
|
self._scale_fn_custom.__dict__.update(fn)
|
|
self._init_scale_fn()
|
|
|
|
|
|
class CosineAnnealingWarmRestarts(LRScheduler):
|
|
r"""Set the learning rate of each parameter group using a cosine annealing schedule.
|
|
|
|
The :math:`\eta_{max}` is set to the initial lr, :math:`T_{cur}`
|
|
is the number of epochs since the last restart and :math:`T_{i}` is the number
|
|
of epochs between two warm restarts in SGDR:
|
|
|
|
.. math::
|
|
\eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 +
|
|
\cos\left(\frac{T_{cur}}{T_{i}}\pi\right)\right)
|
|
|
|
When :math:`T_{cur}=T_{i}`, set :math:`\eta_t = \eta_{min}`.
|
|
When :math:`T_{cur}=0` after restart, set :math:`\eta_t=\eta_{max}`.
|
|
|
|
It has been proposed in
|
|
`SGDR: Stochastic Gradient Descent with Warm Restarts`_.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
T_0 (int): Number of iterations until the first restart.
|
|
T_mult (int, optional): A factor by which :math:`T_{i}` increases after a restart. Default: 1.
|
|
eta_min (float, optional): Minimum learning rate. Default: 0.
|
|
last_epoch (int, optional): The index of the last epoch. Default: -1.
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
.. _SGDR\: Stochastic Gradient Descent with Warm Restarts:
|
|
https://arxiv.org/abs/1608.03983
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
T_0: int,
|
|
T_mult=1,
|
|
eta_min=0.0,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
if T_0 <= 0 or not isinstance(T_0, int):
|
|
raise ValueError(f"Expected positive integer T_0, but got {T_0}")
|
|
if T_mult < 1 or not isinstance(T_mult, int):
|
|
raise ValueError(f"Expected integer T_mult >= 1, but got {T_mult}")
|
|
if not isinstance(eta_min, (float, int)):
|
|
raise ValueError(
|
|
f"Expected float or int eta_min, but got {eta_min} of type {type(eta_min)}"
|
|
)
|
|
self.T_0 = T_0
|
|
self.T_i = T_0
|
|
self.T_mult = T_mult
|
|
self.eta_min = eta_min
|
|
self.T_cur = last_epoch
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def get_lr(self):
|
|
"""Compute the initial learning rate."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
return [
|
|
self.eta_min
|
|
+ (base_lr - self.eta_min)
|
|
* (1 + math.cos(math.pi * self.T_cur / self.T_i))
|
|
/ 2
|
|
for base_lr in self.base_lrs
|
|
]
|
|
|
|
def step(self, epoch=None):
|
|
"""Step could be called after every batch update.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP("Undefined vars")
|
|
>>> scheduler = CosineAnnealingWarmRestarts(optimizer, T_0, T_mult)
|
|
>>> iters = len(dataloader)
|
|
>>> for epoch in range(20):
|
|
>>> for i, sample in enumerate(dataloader):
|
|
>>> inputs, labels = sample['inputs'], sample['labels']
|
|
>>> optimizer.zero_grad()
|
|
>>> outputs = net(inputs)
|
|
>>> loss = criterion(outputs, labels)
|
|
>>> loss.backward()
|
|
>>> optimizer.step()
|
|
>>> scheduler.step(epoch + i / iters)
|
|
|
|
This function can be called in an interleaved way.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP("Undefined vars")
|
|
>>> scheduler = CosineAnnealingWarmRestarts(optimizer, T_0, T_mult)
|
|
>>> for epoch in range(20):
|
|
>>> scheduler.step()
|
|
>>> scheduler.step(26)
|
|
>>> scheduler.step() # scheduler.step(27), instead of scheduler(20)
|
|
"""
|
|
if epoch is None and self.last_epoch < 0:
|
|
epoch = 0
|
|
|
|
if epoch is None:
|
|
epoch = self.last_epoch + 1
|
|
self.T_cur = self.T_cur + 1
|
|
if self.T_cur >= self.T_i:
|
|
self.T_cur = self.T_cur - self.T_i
|
|
self.T_i = self.T_i * self.T_mult
|
|
else:
|
|
if epoch < 0:
|
|
raise ValueError(f"Expected non-negative epoch, but got {epoch}")
|
|
if epoch >= self.T_0:
|
|
if self.T_mult == 1:
|
|
self.T_cur = epoch % self.T_0
|
|
else:
|
|
n = int(
|
|
math.log(
|
|
(epoch / self.T_0 * (self.T_mult - 1) + 1), self.T_mult
|
|
)
|
|
)
|
|
self.T_cur = epoch - self.T_0 * (self.T_mult**n - 1) / (
|
|
self.T_mult - 1
|
|
)
|
|
self.T_i = self.T_0 * self.T_mult ** (n)
|
|
else:
|
|
self.T_i = self.T_0
|
|
self.T_cur = epoch
|
|
self.last_epoch = math.floor(epoch)
|
|
|
|
with _enable_get_lr_call(self):
|
|
for i, data in enumerate(zip(self.optimizer.param_groups, self.get_lr())):
|
|
param_group, lr = data
|
|
param_group["lr"] = lr
|
|
|
|
self._last_lr = [group["lr"] for group in self.optimizer.param_groups]
|
|
|
|
|
|
class _SchedulePhase(TypedDict):
|
|
end_step: float
|
|
start_lr: str
|
|
end_lr: str
|
|
start_momentum: str
|
|
end_momentum: str
|
|
|
|
|
|
class OneCycleLR(LRScheduler):
|
|
r"""Sets the learning rate of each parameter group according to the 1cycle learning rate policy.
|
|
|
|
The 1cycle policy anneals the learning rate from an initial learning rate to some maximum
|
|
learning rate and then from that maximum learning rate to some minimum learning rate much
|
|
lower than the initial learning rate.
|
|
This policy was initially described in the paper `Super-Convergence:
|
|
Very Fast Training of Neural Networks Using Large Learning Rates`_.
|
|
|
|
The 1cycle learning rate policy changes the learning rate after every batch.
|
|
`step` should be called after a batch has been used for training.
|
|
|
|
This scheduler is not chainable.
|
|
|
|
Note also that the total number of steps in the cycle can be determined in one
|
|
of two ways (listed in order of precedence):
|
|
|
|
#. A value for total_steps is explicitly provided.
|
|
#. A number of epochs (epochs) and a number of steps per epoch
|
|
(steps_per_epoch) are provided.
|
|
In this case, the number of total steps is inferred by
|
|
total_steps = epochs * steps_per_epoch
|
|
|
|
You must either provide a value for total_steps or provide a value for both
|
|
epochs and steps_per_epoch.
|
|
|
|
The default behaviour of this scheduler follows the fastai implementation of 1cycle, which
|
|
claims that "unpublished work has shown even better results by using only two phases". To
|
|
mimic the behaviour of the original paper instead, set ``three_phase=True``.
|
|
|
|
Args:
|
|
optimizer (Optimizer): Wrapped optimizer.
|
|
max_lr (float or list): Upper learning rate boundaries in the cycle
|
|
for each parameter group.
|
|
total_steps (int): The total number of steps in the cycle. Note that
|
|
if a value is not provided here, then it must be inferred by providing
|
|
a value for epochs and steps_per_epoch.
|
|
Default: None
|
|
epochs (int): The number of epochs to train for. This is used along
|
|
with steps_per_epoch in order to infer the total number of steps in the cycle
|
|
if a value for total_steps is not provided.
|
|
Default: None
|
|
steps_per_epoch (int): The number of steps per epoch to train for. This is
|
|
used along with epochs in order to infer the total number of steps in the
|
|
cycle if a value for total_steps is not provided.
|
|
Default: None
|
|
pct_start (float): The percentage of the cycle (in number of steps) spent
|
|
increasing the learning rate.
|
|
Default: 0.3
|
|
anneal_strategy (str): {'cos', 'linear'}
|
|
Specifies the annealing strategy: "cos" for cosine annealing, "linear" for
|
|
linear annealing.
|
|
Default: 'cos'
|
|
cycle_momentum (bool): If ``True``, momentum is cycled inversely
|
|
to learning rate between 'base_momentum' and 'max_momentum'.
|
|
Default: True
|
|
base_momentum (float or list): Lower momentum boundaries in the cycle
|
|
for each parameter group. Note that momentum is cycled inversely
|
|
to learning rate; at the peak of a cycle, momentum is
|
|
'base_momentum' and learning rate is 'max_lr'.
|
|
Default: 0.85
|
|
max_momentum (float or list): Upper momentum boundaries in the cycle
|
|
for each parameter group. Functionally,
|
|
it defines the cycle amplitude (max_momentum - base_momentum).
|
|
Note that momentum is cycled inversely
|
|
to learning rate; at the start of a cycle, momentum is 'max_momentum'
|
|
and learning rate is 'base_lr'
|
|
Default: 0.95
|
|
div_factor (float): Determines the initial learning rate via
|
|
initial_lr = max_lr/div_factor
|
|
Default: 25
|
|
final_div_factor (float): Determines the minimum learning rate via
|
|
min_lr = initial_lr/final_div_factor
|
|
Default: 1e4
|
|
three_phase (bool): If ``True``, use a third phase of the schedule to annihilate the
|
|
learning rate according to 'final_div_factor' instead of modifying the second
|
|
phase (the first two phases will be symmetrical about the step indicated by
|
|
'pct_start').
|
|
last_epoch (int): The index of the last batch. This parameter is used when
|
|
resuming a training job. Since `step()` should be invoked after each
|
|
batch instead of after each epoch, this number represents the total
|
|
number of *batches* computed, not the total number of epochs computed.
|
|
When last_epoch=-1, the schedule is started from the beginning.
|
|
Default: -1
|
|
verbose (bool | str): If ``True``, prints a message to stdout for
|
|
each update. Default: ``False``.
|
|
|
|
.. deprecated:: 2.2
|
|
``verbose`` is deprecated. Please use ``get_last_lr()`` to access the
|
|
learning rate.
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> data_loader = torch.utils.data.DataLoader(...)
|
|
>>> optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
|
|
>>> scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.01, steps_per_epoch=len(data_loader), epochs=10)
|
|
>>> for epoch in range(10):
|
|
>>> for batch in data_loader:
|
|
>>> train_batch(...)
|
|
>>> optimizer.step()
|
|
>>> scheduler.step()
|
|
|
|
|
|
.. _Super-Convergence\: Very Fast Training of Neural Networks Using Large Learning Rates:
|
|
https://arxiv.org/abs/1708.07120
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
optimizer: Optimizer,
|
|
max_lr: Union[float, List[float]],
|
|
total_steps: Optional[int] = None,
|
|
epochs: Optional[int] = None,
|
|
steps_per_epoch: Optional[int] = None,
|
|
pct_start=0.3,
|
|
anneal_strategy: Literal["cos", "linear"] = "cos",
|
|
cycle_momentum=True,
|
|
base_momentum: Union[float, List[float]] = 0.85,
|
|
max_momentum: Union[float, List[float]] = 0.95,
|
|
div_factor=25.0,
|
|
final_div_factor=1e4,
|
|
three_phase=False,
|
|
last_epoch=-1,
|
|
verbose="deprecated",
|
|
): # noqa: D107
|
|
# Validate optimizer
|
|
if not isinstance(optimizer, Optimizer):
|
|
raise TypeError(f"{type(optimizer).__name__} is not an Optimizer")
|
|
self.optimizer = optimizer
|
|
|
|
# Validate total_steps
|
|
if total_steps is not None:
|
|
if total_steps <= 0 or not isinstance(total_steps, int):
|
|
raise ValueError(
|
|
f"Expected positive integer total_steps, but got {total_steps}"
|
|
)
|
|
self.total_steps = total_steps
|
|
elif epochs is not None and steps_per_epoch is not None:
|
|
if not isinstance(epochs, int) or epochs <= 0:
|
|
raise ValueError(f"Expected positive integer epochs, but got {epochs}")
|
|
if not isinstance(steps_per_epoch, int) or steps_per_epoch <= 0:
|
|
raise ValueError(
|
|
f"Expected positive integer steps_per_epoch, but got {steps_per_epoch}"
|
|
)
|
|
self.total_steps = epochs * steps_per_epoch
|
|
else:
|
|
raise ValueError(
|
|
"You must define either total_steps OR (epochs AND steps_per_epoch)"
|
|
)
|
|
|
|
self._schedule_phases: List[_SchedulePhase]
|
|
if three_phase:
|
|
self._schedule_phases = [
|
|
{
|
|
"end_step": float(pct_start * self.total_steps) - 1,
|
|
"start_lr": "initial_lr",
|
|
"end_lr": "max_lr",
|
|
"start_momentum": "max_momentum",
|
|
"end_momentum": "base_momentum",
|
|
},
|
|
{
|
|
"end_step": float(2 * pct_start * self.total_steps) - 2,
|
|
"start_lr": "max_lr",
|
|
"end_lr": "initial_lr",
|
|
"start_momentum": "base_momentum",
|
|
"end_momentum": "max_momentum",
|
|
},
|
|
{
|
|
"end_step": self.total_steps - 1,
|
|
"start_lr": "initial_lr",
|
|
"end_lr": "min_lr",
|
|
"start_momentum": "max_momentum",
|
|
"end_momentum": "max_momentum",
|
|
},
|
|
]
|
|
else:
|
|
self._schedule_phases = [
|
|
{
|
|
"end_step": float(pct_start * self.total_steps) - 1,
|
|
"start_lr": "initial_lr",
|
|
"end_lr": "max_lr",
|
|
"start_momentum": "max_momentum",
|
|
"end_momentum": "base_momentum",
|
|
},
|
|
{
|
|
"end_step": self.total_steps - 1,
|
|
"start_lr": "max_lr",
|
|
"end_lr": "min_lr",
|
|
"start_momentum": "base_momentum",
|
|
"end_momentum": "max_momentum",
|
|
},
|
|
]
|
|
|
|
# Validate pct_start
|
|
if pct_start < 0 or pct_start > 1 or not isinstance(pct_start, float):
|
|
raise ValueError(
|
|
f"Expected float between 0 and 1 pct_start, but got {pct_start}"
|
|
)
|
|
|
|
# Validate anneal_strategy
|
|
if anneal_strategy not in ["cos", "linear"]:
|
|
raise ValueError(
|
|
f"anneal_strategy must be one of 'cos' or 'linear', instead got {anneal_strategy}"
|
|
)
|
|
else:
|
|
self._anneal_func_type = anneal_strategy
|
|
|
|
# Initialize learning rate variables
|
|
max_lrs = _format_param("max_lr", self.optimizer, max_lr)
|
|
if last_epoch == -1:
|
|
for idx, group in enumerate(self.optimizer.param_groups):
|
|
group["initial_lr"] = max_lrs[idx] / div_factor
|
|
group["max_lr"] = max_lrs[idx]
|
|
group["min_lr"] = group["initial_lr"] / final_div_factor
|
|
|
|
# Initialize momentum variables
|
|
self.cycle_momentum = cycle_momentum
|
|
if self.cycle_momentum:
|
|
if (
|
|
"momentum" not in self.optimizer.defaults
|
|
and "betas" not in self.optimizer.defaults
|
|
):
|
|
raise ValueError(
|
|
"optimizer must support momentum or beta1 with `cycle_momentum` option enabled"
|
|
)
|
|
self.use_beta1 = "betas" in self.optimizer.defaults
|
|
max_momentums = _format_param("max_momentum", optimizer, max_momentum)
|
|
base_momentums = _format_param("base_momentum", optimizer, base_momentum)
|
|
if last_epoch == -1:
|
|
for m_momentum, b_momentum, group in zip(
|
|
max_momentums, base_momentums, optimizer.param_groups
|
|
):
|
|
if self.use_beta1:
|
|
group["betas"] = (m_momentum, *group["betas"][1:])
|
|
else:
|
|
group["momentum"] = m_momentum
|
|
group["max_momentum"] = m_momentum
|
|
group["base_momentum"] = b_momentum
|
|
|
|
super().__init__(optimizer, last_epoch, verbose)
|
|
|
|
def _anneal_func(self, *args, **kwargs):
|
|
if hasattr(self, "_anneal_func_type"):
|
|
if self._anneal_func_type == "cos":
|
|
return self._annealing_cos(*args, **kwargs)
|
|
elif self._anneal_func_type == "linear":
|
|
return self._annealing_linear(*args, **kwargs)
|
|
else:
|
|
raise ValueError(f"Unknown _anneal_func_type: {self._anneal_func_type}")
|
|
else:
|
|
# For BC
|
|
return self.anneal_func(*args, **kwargs) # type: ignore[attr-defined]
|
|
|
|
@staticmethod
|
|
def _annealing_cos(start, end, pct):
|
|
"""Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0."""
|
|
cos_out = math.cos(math.pi * pct) + 1
|
|
return end + (start - end) / 2.0 * cos_out
|
|
|
|
@staticmethod
|
|
def _annealing_linear(start, end, pct):
|
|
"""Linearly anneal from `start` to `end` as pct goes from 0.0 to 1.0."""
|
|
return (end - start) * pct + start
|
|
|
|
def get_lr(self):
|
|
"""Compute the learning rate of each parameter group."""
|
|
_warn_get_lr_called_within_step(self)
|
|
|
|
lrs = []
|
|
step_num = self.last_epoch
|
|
|
|
if step_num > self.total_steps:
|
|
raise ValueError(
|
|
f"Tried to step {step_num} times. The specified number of total steps is {self.total_steps}" # noqa: UP032
|
|
)
|
|
|
|
for group in self.optimizer.param_groups:
|
|
start_step = 0.0
|
|
for i, phase in enumerate(self._schedule_phases):
|
|
end_step = phase["end_step"]
|
|
if step_num <= end_step or i == len(self._schedule_phases) - 1:
|
|
pct = (step_num - start_step) / (end_step - start_step)
|
|
computed_lr = self._anneal_func(
|
|
group[phase["start_lr"]], group[phase["end_lr"]], pct
|
|
)
|
|
if self.cycle_momentum:
|
|
computed_momentum = self._anneal_func(
|
|
group[phase["start_momentum"]],
|
|
group[phase["end_momentum"]],
|
|
pct,
|
|
)
|
|
break
|
|
start_step = phase["end_step"]
|
|
|
|
lrs.append(computed_lr) # type: ignore[possibly-undefined]
|
|
if self.cycle_momentum:
|
|
if self.use_beta1:
|
|
group["betas"] = (computed_momentum, *group["betas"][1:]) # type: ignore[possibly-undefined]
|
|
else:
|
|
group[
|
|
"momentum"
|
|
] = computed_momentum # type: ignore[possibly-undefined]
|
|
|
|
return lrs
|