# coding=utf-8
__author__ = "Dimitris Karkalousos"
# Taken and adapted from: https://github.com/NVIDIA/NeMo/blob/main/nemo/core/optim/lr_scheduler.py
import copy
import dataclasses
import inspect
import math
import warnings
from functools import partial
from typing import Any, Dict, Optional, Union
import hydra
import torch.optim.lr_scheduler as pt_scheduler
from omegaconf import DictConfig, OmegaConf
from torch import optim
from torch.optim.lr_scheduler import _LRScheduler
from torch.utils.data import dataloader
from atommic.core.conf.schedulers import SchedulerParams, get_scheduler_config, register_scheduler_params
from atommic.utils import logging
from atommic.utils.model_utils import maybe_update_config_version
class WarmupPolicy(_LRScheduler):
"""Adds warmup kwargs and warmup logic to lr policy. All arguments should be passed as kwargs for clarity.
Returns
-------
lr : float
Learning rate for current step.
"""
def __init__(self, optimizer, *, warmup_steps=None, warmup_ratio=None, max_steps=None, min_lr=0.0, last_epoch=-1):
"""Inits :class:`WarmupPolicy`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer.
warmup_steps : int
Number of training steps in warmup stage. Default is ``None``.
warmup_ratio : float
Ratio of warmup steps to total steps. Default is ``None``.
max_steps : int
Total number of steps while training or `None` for infinite training. Default is ``None``.
min_lr : float
Minimum learning rate. Default is ``0``.
last_epoch : int
Last epoch. Default is ``-1``.
"""
if warmup_steps is not None and warmup_ratio is not None:
raise AssertionError("Either use particular number of step or ratio")
if warmup_ratio is not None and max_steps is None:
raise AssertionError("If there is a ratio, there should be a total steps")
# It is necessary to assign all attributes *before* __init__,
# as class is wrapped by an inner class.
self.max_steps = max_steps
if warmup_steps is not None:
self.warmup_steps = warmup_steps
elif warmup_ratio is not None:
self.warmup_steps = int(warmup_ratio * max_steps)
else:
self.warmup_steps = 0
self.min_lr = min_lr
super().__init__(optimizer, last_epoch)
def get_lr(self):
"""Get learning rate at current step."""
if not self._get_lr_called_within_step:
warnings.warn(
"To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
)
step = self.last_epoch
if step <= self.warmup_steps and self.warmup_steps > 0:
return self._get_warmup_lr(step)
if (self.max_steps is not None) and (step > self.max_steps):
return [self.min_lr for _ in self.base_lrs]
return self._get_lr(step)
def _get_warmup_lr(self, step):
"""Linear warmup"""
lr_val = (step + 1) / (self.warmup_steps + 1)
return [initial_lr * lr_val for initial_lr in self.base_lrs]
def _get_lr(self, step): # pylint: disable=unused-argument
"""Simple const lr policy"""
return self.base_lrs
class SquareRootConstantPolicy(_LRScheduler):
"""Adds warmup kwargs and warmup logic to lr policy. All arguments should be passed as kwargs for clarity."""
def __init__(
self, optimizer, *, constant_steps=None, constant_ratio=None, max_steps=None, min_lr=0.0, last_epoch=-1
):
"""Inits :class:`SquareRootConstantPolicy`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer.
constant_steps : int
Number of training steps in constant stage. Default is ``None``.
constant_ratio : float
Ratio of constant steps to total steps. Default is ``None``.
max_steps : int
Total number of steps while training or `None` for infinite training. Default is ``None``.
min_lr : float
Minimum learning rate. Default is ``0``.
last_epoch : int
Last epoch. Default is ``-1``.
"""
if constant_steps is not None and constant_ratio is not None:
raise AssertionError("Either use particular number of step or ratio")
if constant_ratio is not None and max_steps is None:
raise AssertionError("If there is a ratio, there should be a total steps")
# It is necessary to assign all attributes *before* __init__, as class is wrapped by an inner class.
self.max_steps = max_steps
if constant_steps is not None:
self.constant_steps = constant_steps
elif constant_ratio is not None:
self.constant_steps = int(constant_ratio * max_steps)
else:
self.constant_steps = 0
self.constant_lr = 1 / (constant_steps**0.5)
self.min_lr = min_lr
super().__init__(optimizer, last_epoch)
def get_lr(self):
"""Get learning rate at current step."""
if not self._get_lr_called_within_step:
warnings.warn(
"To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
)
step = self.last_epoch
if step <= self.constant_steps:
return [self.constant_lr for _ in self.base_lrs]
if step > self.max_steps:
return [self.min_lr for _ in self.base_lrs]
return self._get_lr(step)
def _get_lr(self, step): # pylint: disable=unused-argument
"""Simple const lr policy"""
return self.base_lrs
class WarmupHoldPolicy(WarmupPolicy):
"""Variant of WarmupPolicy which maintains high learning rate for a defined number of steps. All arguments should
be passed as kwargs for clarity,
Results
-------
lr : float
Learning rate is linearly increased from 0 to 1 over warmup steps, then linearly decreased from 1 to 0 over
hold steps.
"""
def __init__(
self,
optimizer,
*,
warmup_steps=None,
warmup_ratio=None,
hold_steps=None,
hold_ratio=None,
max_steps=None,
min_lr=0.0,
last_epoch=-1,
):
"""Inits :class:`WarmupHoldPolicy`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer.
warmup_steps : int
Number of training steps in warmup stage. Default is ``None``.
warmup_ratio : float
Ratio of warmup steps to total steps. Default is ``None``.
hold_steps : int
Number of training steps to hold the learning rate after warm up. Default is ``None``.
hold_ratio : float
Ratio of hold steps to total steps. Default is ``None``.
max_steps : int
Total number of steps while training or `None` for infinite training. Default is ``None``.
min_lr : float
Minimum learning rate. Default is ``0``.
last_epoch : int
Last epoch. Default is ``-1``.
"""
if hold_steps is not None and hold_ratio is not None:
raise AssertionError("Either use particular number of step or ratio")
if hold_ratio is not None and max_steps is None:
raise AssertionError("If there is a ratio, there should be a total steps")
self.min_lr = min_lr
self._last_warmup_lr = 0.0
# Necessary to duplicate as class attributes are hidden in inner class
self.max_steps = max_steps
if warmup_steps is not None:
self.warmup_steps = warmup_steps
elif warmup_ratio is not None:
self.warmup_steps = int(warmup_ratio * max_steps)
else:
self.warmup_steps = 0
if hold_steps is not None:
self.hold_steps = hold_steps + self.warmup_steps
elif hold_ratio is not None:
self.hold_steps = int(hold_ratio * max_steps) + self.warmup_steps
else:
self.hold_steps = 0
super().__init__(
optimizer,
warmup_steps=warmup_steps,
warmup_ratio=warmup_ratio,
max_steps=max_steps,
last_epoch=last_epoch,
min_lr=min_lr,
)
def get_lr(self):
"""Get learning rate at current step."""
if not self._get_lr_called_within_step:
warnings.warn(
"To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
)
step = self.last_epoch
# Warmup phase
if 0 < self.warmup_steps >= step:
return self._get_warmup_lr(step)
# Hold phase
if self.hold_steps < step >= self.warmup_steps:
return self.base_lrs
if step > self.max_steps:
return [self.min_lr for _ in self.base_lrs]
return self._get_lr(step)
class WarmupAnnealHoldPolicy(_LRScheduler):
"""Adds warmup kwargs and warmup logic to lr policy. All arguments should be passed as kwargs for clarity."""
def __init__(
self,
optimizer,
*,
warmup_steps=None,
warmup_ratio=None,
constant_steps=None,
constant_ratio=None,
max_steps=None,
min_lr=0.0,
last_epoch=-1,
):
"""Inits :class:`WarmupAnnealHoldPolicy`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer.
warmup_steps : int
Number of training steps in warmup stage. Default is ``None``.
warmup_ratio : float
Ratio of warmup steps to total steps. Default is ``None``.
constant_steps : int
Number of training steps in constant stage. Default is ``None``.
constant_ratio : float
Ratio of constant steps to total steps. Default is ``None``.
max_steps : int
Total number of steps while training or `None` for infinite training. Default is ``None``.
min_lr : float
Minimum learning rate. Default is ``0``.
last_epoch : int
Last epoch. Default is ``-1``.
"""
if warmup_steps is not None and warmup_ratio is not None:
raise AssertionError("Either use particular number of step or ratio")
if constant_steps is not None and constant_ratio is not None:
raise AssertionError("Either use constant_steps or constant_ratio")
if warmup_ratio is not None and max_steps is None:
raise AssertionError("If there is a ratio, there should be a total steps")
# It is necessary to assign all attributes *before* __init__, as class is wrapped by an inner class.
self.max_steps = max_steps
if warmup_steps is not None:
self.warmup_steps = warmup_steps
elif warmup_ratio is not None:
self.warmup_steps = int(warmup_ratio * max_steps)
else:
self.warmup_steps = 0
if constant_steps is not None:
self.constant_steps = constant_steps
elif constant_ratio is not None:
self.constant_steps = int(constant_ratio * max_steps)
else:
self.constant_steps = 0
self.decay_steps = max_steps - (self.constant_steps + self.warmup_steps)
self.min_lr = min_lr
super().__init__(optimizer, last_epoch)
def get_lr(self):
"""Get learning rate at current step."""
if not self._get_lr_called_within_step:
warnings.warn(
"To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
)
step = self.last_epoch
# Warmup steps
if 0 < self.warmup_steps >= step:
return self._get_warmup_lr(step)
# Constant steps after warmup and decay
if self.constant_steps > 0 and (self.warmup_steps + self.decay_steps) < step <= self.max_steps:
return self._get_constant_lr(step)
# Min lr after max steps of updates
if step > self.max_steps:
return [self.min_lr for _ in self.base_lrs]
return self._get_lr(step)
def _get_warmup_lr(self, step):
"""Get learning rate at warmup stage."""
lr_val = (step + 1) / (self.warmup_steps + 1)
return [initial_lr * lr_val for initial_lr in self.base_lrs]
def _get_constant_lr(self, step): # pylint: disable=unused-argument
"""Get learning rate at constant stage."""
return [self.min_lr for _ in self.base_lrs]
def _get_lr(self, step): # pylint: disable=unused-argument
"""Simple const lr policy"""
return self.base_lrs
def _sqrt_annealing(initial_lr, step, max_steps, min_lr):
"""Anneal learning rate by sqrt."""
mult = ((max_steps - step) / max_steps) ** 0.5
out_lr = initial_lr * mult
out_lr = max(out_lr, min_lr)
return out_lr
def _square_annealing(initial_lr, step, max_steps, min_lr):
"""Anneal learning rate by square."""
mult = ((max_steps - step) / max_steps) ** 2
out_lr = initial_lr * mult
out_lr = max(out_lr, min_lr)
return out_lr
def _cosine_annealing(initial_lr, step, max_steps, min_lr):
"""Anneal learning rate by cosine."""
mult = 0.5 * (1 + math.cos(math.pi * step / max_steps))
return (initial_lr - min_lr) * mult + min_lr
def _linear_warmup_with_cosine_annealing(max_lr, warmup_steps, step, decay_steps, min_lr):
"""Anneal learning rate by linear warmup and cosine annealing."""
if max_lr <= min_lr:
raise AssertionError
# Use linear warmup for the initial part.
if warmup_steps > 0 and step <= warmup_steps:
return max_lr * float(step) / float(warmup_steps)
# For any steps larger than `decay_steps`, use `min_lr`.
if step > warmup_steps + decay_steps:
return min_lr
# If we are done with the warmup period, use the decay style.
num_steps_ = step - warmup_steps
decay_steps_ = decay_steps
decay_ratio = float(num_steps_) / float(decay_steps_)
if decay_ratio < 0.0:
raise AssertionError
if decay_ratio > 1.0:
raise AssertionError
delta_lr = max_lr - min_lr
coeff = 0.5 * (math.cos(math.pi * decay_ratio) + 1.0)
return min_lr + coeff * delta_lr
def _poly_decay(initial_lr, step, decay_steps, power, min_lr, cycle):
"""Polynomial decay of learning rate."""
if cycle:
multiplier = 1.0 if step == 0 else math.ceil(step / decay_steps)
decay_steps = decay_steps * multiplier
else:
step = min(step, decay_steps)
p = step / decay_steps
lr = (initial_lr - min_lr) * math.pow(1.0 - p, power)
lr += min_lr
return lr
def _noam_hold_annealing(initial_lr, step, warmup_steps, hold_steps, decay_rate, min_lr):
"""Anneal learning rate by noam hold."""
# hold_steps = total number of steps to hold the LR, not the warmup + hold steps.
T_warmup_decay = max(1, warmup_steps**decay_rate)
T_hold_decay = max(1, (step - hold_steps) ** decay_rate)
lr = (initial_lr * T_warmup_decay) / T_hold_decay
return max(lr, min_lr)
class SquareAnnealing(WarmupPolicy):
"""Anneal learning rate by square."""
def __init__(self, optimizer, *, max_steps, min_lr=1e-5, last_epoch=-1, **kwargs):
"""Inits :class:`SquareAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer.
max_steps : int
Total number of steps while training or `None` for infinite training.
min_lr : float
Minimum learning rate. Default is ``1e-5``.
last_epoch : int
Last epoch. Default is ``-1``.
"""
super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
def _get_lr(self, step):
"""Get learning rate at current step."""
return [
_square_annealing(
initial_lr=initial_lr,
step=step - self.warmup_steps,
max_steps=self.max_steps - self.warmup_steps,
min_lr=self.min_lr,
)
for initial_lr in self.base_lrs
]
class SquareRootAnnealing(WarmupPolicy):
"""Anneal learning rate by square root."""
def __init__(self, optimizer, *, max_steps, min_lr=0, last_epoch=-1, **kwargs):
"""Inits :class:`SquareRootAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer.
max_steps : int
Total number of steps while training or `None` for infinite training.
min_lr : float
Minimum learning rate. Default is ``0``.
last_epoch : int
Last epoch. Default is ``-1``.
"""
super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
def _get_lr(self, step):
"""Get learning rate at current step."""
return [
_sqrt_annealing(
initial_lr=initial_lr,
step=step,
max_steps=self.max_steps,
min_lr=self.min_lr,
)
for initial_lr in self.base_lrs
]
class CosineAnnealing(WarmupAnnealHoldPolicy):
"""Anneal learning rate by cosine."""
def __init__(self, optimizer, *, max_steps, min_lr=0, last_epoch=-1, **kwargs):
"""Inits :class:`CosineAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer.
max_steps : int
Total number of steps while training or `None` for infinite training.
min_lr : float
Minimum learning rate. Default is ``0``.
last_epoch : int
Last epoch. Default is ``-1``.
"""
super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
def _get_lr(self, step):
"""Get learning rate at current step."""
for initial_lr in self.base_lrs:
if initial_lr < self.min_lr:
raise ValueError(
f"{self} received an initial learning rate that was lower than the minimum learning rate."
)
return (
[
_cosine_annealing(
initial_lr=initial_lr,
step=step - self.warmup_steps,
max_steps=self.max_steps - self.warmup_steps,
min_lr=self.min_lr,
)
for initial_lr in self.base_lrs
]
if self.constant_steps is None or self.constant_steps == 0
else self._get_linear_warmup_with_cosine_annealing_lr(step)
)
def _get_warmup_lr(self, step):
"""Get the warmup learning rate for the given step."""
if self.constant_steps is None or self.constant_steps == 0:
return super()._get_warmup_lr(step)
# Use linear warmup for the initial part.
return self._get_linear_warmup_with_cosine_annealing_lr(step)
def _get_constant_lr(self, step):
"""Only called when constant_steps is not None and not 0."""
return self._get_linear_warmup_with_cosine_annealing_lr(step)
def _get_linear_warmup_with_cosine_annealing_lr(self, step):
"""Cosine Schedule, slightly different warmup schedule + constant LR at the end."""
return [
_linear_warmup_with_cosine_annealing(
max_lr=self.base_lrs[0],
warmup_steps=self.warmup_steps,
step=step,
decay_steps=self.decay_steps,
min_lr=self.min_lr,
)
for _ in self.base_lrs
]
class NoamAnnealing(_LRScheduler):
"""Noam learning rate annealing."""
def __init__(
self, optimizer, *, d_model, warmup_steps=None, warmup_ratio=None, max_steps=None, min_lr=0.0, last_epoch=-1
):
"""Inits :class:`NoamAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer.
d_model : int
Model dimensionality.
warmup_steps : int
Number of training steps in warmup stage. Default is ``None``.
warmup_ratio : float
Ratio of warmup steps to total steps. Default is ``None``.
max_steps : int
Total number of steps while training or `None` for infinite training.
min_lr : float
Minimum learning rate. Default is ``0``.
last_epoch : int
Last epoch. Default is ``-1``.
"""
self._normalize = d_model ** (-0.5)
if warmup_steps is not None and warmup_ratio is not None:
raise AssertionError("Either use particular number of step or ratio")
if warmup_ratio is not None and max_steps is None:
raise AssertionError("If there is a ratio, there should be a total steps")
# It is necessary to assign all attributes *before* __init__,
# as class is wrapped by an inner class.
self.max_steps = max_steps
if warmup_steps is not None:
self.warmup_steps = warmup_steps
elif warmup_ratio is not None:
self.warmup_steps = int(warmup_ratio * max_steps)
else:
self.warmup_steps = 0
self.min_lr = min_lr
super().__init__(optimizer, last_epoch)
def get_lr(self):
"""Get learning rate at current step."""
if not self._get_lr_called_within_step:
warnings.warn(
"To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
)
step = max(1, self.last_epoch)
if step > self.max_steps:
return [self.min_lr for _ in self.base_lrs]
for initial_lr in self.base_lrs:
if initial_lr < self.min_lr:
raise ValueError(
f"{self} received an initial learning rate that was lower than the minimum learning rate."
)
return [self._noam_annealing(initial_lr=initial_lr, step=step) for initial_lr in self.base_lrs]
def _noam_annealing(self, initial_lr, step):
"""Noam learning rate annealing."""
mult = (
self._normalize * min(step ** (-0.5), step * (self.warmup_steps ** (-1.5)))
if self.warmup_steps > 0
else self._normalize * step ** (-0.5)
)
out_lr = initial_lr * mult
if step > self.warmup_steps:
out_lr = max(out_lr, self.min_lr)
return out_lr
class NoamHoldAnnealing(WarmupHoldPolicy):
"""Implementation of the Noam Hold Annealing policy from the SqueezeFormer paper.
Unlike NoamAnnealing, the peak learning rate can be explicitly set for this scheduler.
The schedule first performs linear warmup, then holds the peak LR, then decays with some schedule for
the remainder of the steps. Therefore, the min-lr is still dependent on the hyperparameters selected.
It's schedule is determined by three factors-
Warmup Steps: Initial stage, where linear warmup occurs uptil the peak LR is reached. Unlike NoamAnnealing,
the peak LR is explicitly stated here instead of a scaling factor.
Hold Steps: Intermediate stage, where the peak LR is maintained for some number of steps. In this region,
the high peak LR allows the model to converge faster if training is stable. However the high LR
may also cause instability during training. Should usually be a significant fraction of training
steps (around 30-40% of the entire training steps).
Decay Steps: Final stage, where the LR rapidly decays with some scaling rate (set by decay rate).
To attain Noam decay, use 0.5, for Squeezeformer recommended decay, use 1.0. The fast decay after
prolonged high LR during hold phase allows for rapid convergence.
References
----------
[1]
[Squeezeformer: An Efficient Transformer for Automatic Speech Recognition](https://arxiv.org/abs/2206.00888)
"""
def __init__(self, optimizer, *, max_steps, decay_rate=0.5, min_lr=0.0, last_epoch=-1, **kwargs):
"""Inits :class:`NoamHoldAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer to use for the scheduler.
max_steps : int
Total number of training steps.
decay_rate : float
Decay rate for the final stage of the schedule. Should be between 0 and 1. Default is ``0.5``.
min_lr : float
Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
last_epoch : int
Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
"""
self.decay_rate = decay_rate
super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
def _get_lr(self, step):
"""Get the learning rate for the given step."""
if self.warmup_steps is None or self.warmup_steps == 0:
raise ValueError("Noam scheduler cannot be used without warmup steps")
if self.hold_steps > 0:
hold_steps = self.hold_steps - self.warmup_steps
else:
hold_steps = 0
return [
_noam_hold_annealing(
initial_lr,
step=step,
warmup_steps=self.warmup_steps,
hold_steps=hold_steps,
decay_rate=self.decay_rate,
min_lr=self.min_lr,
)
for initial_lr in self.base_lrs
]
class WarmupAnnealing(WarmupPolicy):
"""Warmup learning rate annealing."""
def __init__(self, optimizer, *, max_steps, last_epoch=-1, min_lr=0.0, **kwargs):
"""Inits :class:`WarmupAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer to use for the scheduler.
max_steps : int
Total number of training steps.
last_epoch : int
Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
min_lr : float
Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
"""
super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
def _get_lr(self, step):
"""Get learning rate at current step."""
delta_lr = self.base_lrs[0] - self.min_lr
mult = (step - self.warmup_steps) / (self.max_steps - self.warmup_steps)
return [self.min_lr + (1 - mult) * delta_lr for _ in self.base_lrs]
class InverseSquareRootAnnealing(WarmupPolicy):
"""Inverse square root learning rate annealing."""
def __init__(self, optimizer, *, max_steps, last_epoch=-1, min_lr=0.0, **kwargs):
"""Inits :class:`InverseSquareRootAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer to use for the scheduler.
max_steps : int
Total number of training steps.
last_epoch : int
Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
min_lr : float
Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
"""
super().__init__(optimizer=optimizer, max_steps=max_steps, **kwargs, last_epoch=last_epoch, min_lr=min_lr)
def _get_lr(self, step):
"""Get learning rate at current step."""
denom = ((step + 1) / (self.warmup_steps + 1)) ** 0.5
return [initial_lr / denom for initial_lr in self.base_lrs]
class T5InverseSquareRootAnnealing(SquareRootConstantPolicy):
"""Inverse square root learning rate annealing."""
def __init__(self, optimizer, *, max_steps, last_epoch=-1, min_lr=0.0, **kwargs):
"""Inits :class:`T5InverseSquareRootAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer to use for the scheduler.
max_steps : int
Total number of training steps.
last_epoch : int
Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
min_lr : float
Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
"""
super().__init__(optimizer=optimizer, max_steps=max_steps, **kwargs, last_epoch=last_epoch, min_lr=min_lr)
def _get_lr(self, step):
"""Get learning rate at current step."""
return [1 / (step**0.5) for _ in self.base_lrs]
class PolynomialDecayAnnealing(WarmupPolicy):
"""Polynomial decay learning rate annealing."""
def __init__(self, optimizer, *, max_steps, min_lr=0.0, power=1.0, cycle=False, last_epoch=-1, **kwargs):
"""Inits :class:`PolynomialDecayAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer to use for the scheduler.
max_steps : int
Total number of training steps.
min_lr : float
Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
power : float
Power of the polynomial. Default is ``1.0``.
cycle : bool
Whether to cycle the schedule. Default is ``False``.
last_epoch : int
Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
"""
self.power = power
self.cycle = cycle
super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
def _get_lr(self, step):
"""Get learning rate at current step."""
return [
_poly_decay(
initial_lr,
step=step - self.warmup_steps,
decay_steps=self.max_steps - self.warmup_steps,
power=self.power,
min_lr=self.min_lr,
cycle=self.cycle,
)
for initial_lr in self.base_lrs
]
class PolynomialHoldDecayAnnealing(WarmupHoldPolicy):
"""Polynomial decay learning rate annealing."""
def __init__(self, optimizer, *, max_steps, min_lr=0.0, power=1.0, cycle=False, last_epoch=-1, **kwargs):
"""Inits :class:`PolynomialHoldDecayAnnealing`.
Parameters
----------
optimizer : torch.optim.Optimizer
Optimizer to use for the scheduler.
max_steps : int
Total number of training steps.
min_lr : float
Minimum learning rate to use for the schedule. Should be between 0 and 1. Default is ``0.0``.
power : float
Power of the polynomial. Default is ``1.0``.
cycle : bool
Whether to cycle the schedule. Default is ``False``.
last_epoch : int
Last epoch to start the schedule from. Should be between 0 and max_steps. Default is ``-1``.
"""
self.power = power
self.cycle = cycle
super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
def _get_lr(self, step):
"""Get learning rate at current step."""
return [
_poly_decay(
initial_lr,
step=step - self.hold_steps,
decay_steps=self.max_steps - max(self.warmup_steps, self.hold_steps),
power=self.power,
min_lr=self.min_lr,
cycle=self.cycle,
)
for initial_lr in self.base_lrs
]
[docs]def register_scheduler(name: str, scheduler: _LRScheduler, scheduler_params: SchedulerParams):
"""Checks if the scheduler name exists in the registry, and if it doesn't, adds it. This allows custom schedulers
to be added and called by name during instantiation.
Parameters
----------
name : str
Name of the optimizer. Will be used as key to retrieve the optimizer.
scheduler : _LRScheduler
Scheduler class (inherits from _LRScheduler)
scheduler_params : SchedulerParams
The parameters as a dataclass of the scheduler
"""
if name in AVAILABLE_SCHEDULERS:
raise ValueError(f"Cannot override pre-existing schedulers. Conflicting scheduler name = {name}")
AVAILABLE_SCHEDULERS[name] = scheduler
sched_name = f"{scheduler.__name__}_params"
register_scheduler_params(name=sched_name, scheduler_params=scheduler_params)
def get_scheduler(name: str, **kwargs: Optional[Dict[str, Any]]) -> _LRScheduler:
"""Convenience method to obtain an _LRScheduler class and partially instantiate it with optimizer kwargs.
Parameters
----------
name : str
Name of the scheduler in the registry.
kwargs : Dict[str, Any]
Optional kwargs of the scheduler used during instantiation.
Returns
-------
_LRScheduler
A partially instantiated _LRScheduler.
"""
if name not in AVAILABLE_SCHEDULERS:
raise ValueError(
f"Cannot resolve scheduler{name}'. Available optimizers are : " f"{AVAILABLE_SCHEDULERS.keys()}"
)
scheduler_cls = AVAILABLE_SCHEDULERS[name]
# Pop 'max_steps' if it's not required by the scheduler
if "max_steps" in kwargs and "max_steps" not in inspect.signature(scheduler_cls).parameters:
kwargs.pop("max_steps")
scheduler = partial(scheduler_cls, **kwargs)
return scheduler
def prepare_lr_scheduler( # noqa: MC0001
optimizer: optim.Optimizer,
scheduler_config: Union[Dict[str, Any], DictConfig, None],
train_dataloader: Optional[dataloader.DataLoader] = None,
) -> Optional[Dict[str, Any]]:
"""Constructs an LR Scheduler (optionally) for a given optimizer, based on a config with the following schema.
Parameters
----------
optimizer: The optimizer to use for the scheduler.
name: <name of optimizer>
lr: <maximal learning rate>
args:
name: auto # special keyword, resolves to correct optimizer config for given optimizer name
# cls: atommic.core.config.optimizers.NovogradParams # explicit instantiation by class path
params: # optional override parameters for the optimizer config
betas: [0.8, 0.5]
weight_decay: 0.001
scheduler_config: The scheduler config.
name: <name of scheduler>
iters_per_batch: null # computed at runtime; mandatory to have
max_steps: null # computed at runtime or explicitly set here; mandatory to have
# pytorch lightning args <mandatory>
monitor: val_loss
reduce_on_plateau: false
# <scheduler config override>
args:
name: auto # special keyword, resolves to correct optimizer config for given optimizer name
# cls: atommic.core.config.schedulers.CosineAnnealingParams # explicit instantiation by class path
params: # optional override parameters for the optimizer config
warmup_steps: null
warmup_ratio: null
min_lr: 0.0
last_epoch: -1
train_dataloader: Optional requirement, must be passed if "iters_per_batch" is defined instead of "max_steps". \
Used to compute effective "max_steps".
Returns
-------
A dictionary containing the LR Scheduler implementation if the config was successfully parsed along with other \
parameters required by Pytorch Lightning, otherwise None.
"""
if scheduler_config is not None:
scheduler_config = maybe_update_config_version(scheduler_config)
# Build nested dictionary for convenience out of structured objects
if isinstance(scheduler_config, DictConfig):
scheduler_config = OmegaConf.to_container(scheduler_config, resolve=True)
elif dataclasses.is_dataclass(scheduler_config):
# Recursively transform data classes to basic dictionaries
scheduler_config = OmegaConf.create(scheduler_config)
scheduler_config = OmegaConf.to_container(scheduler_config, resolve=True)
# Test to see if config follows above schema
interval = "step"
if scheduler_config is not None:
if "args" in scheduler_config:
scheduler_args = scheduler_config.pop("args")
else:
scheduler_args = copy.deepcopy(scheduler_config)
# Remove extra parameters from scheduler_args nest
# Assume all other parameters are to be passed into scheduler constructor
scheduler_args.pop("name", None)
scheduler_args.pop("t_max_epochs", None)
scheduler_args.pop("t_accumulate_grad_batches", None)
scheduler_args.pop("t_limit_train_batches", None)
scheduler_args.pop("t_num_workers", None)
scheduler_args.pop("monitor", None)
scheduler_args.pop("reduce_on_plateau", None)
if "name" in scheduler_config and scheduler_config["name"] in EPOCH_SCHEDULERS:
interval = "epoch"
else:
# Return gracefully in case `sched` was not supplied; inform user
logging.info("Scheduler not initialized as no `sched` config supplied to setup_optimizer()")
return None
# Try instantiation of scheduler params from config class path
if "_target_" in scheduler_args:
scheduler_args_cfg = OmegaConf.create(scheduler_args)
scheduler_conf = hydra.utils.instantiate(scheduler_args_cfg)
scheduler_args = vars(scheduler_conf)
# Get name of the scheduler
scheduler_name = scheduler_conf.__class__.__name__
if "Params" in scheduler_name:
scheduler_name = scheduler_name.replace("Params", "")
else:
# Class path instantiation failed; try resolving "name" component
# Get name of the scheduler
if "name" in scheduler_config:
scheduler_name = scheduler_config["name"]
else:
logging.warning(
"Could not resolve classpath for Scheduler Config, and `name` "
"was not provided either. \n"
"Scheduler cannot be instantiated !"
)
return None
# If class path was not provided, perhaps `name` is provided for resolution
if "name" in scheduler_args:
# If `auto` is passed as name for resolution of optimizer name,
# then lookup optimizer name and resolve its parameter config
if scheduler_args["name"] == "auto":
scheduler_params_name = f"{scheduler_name}Params"
else:
scheduler_params_name = scheduler_args["name"]
# Get override arguments provided in the config yaml file / Dict Config
scheduler_params_override = scheduler_args.get("params", {})
# If params is itself a dict config object provided explicitly in Dict Config
# Resolve to dictionary for convenience
if isinstance(scheduler_params_override, DictConfig):
scheduler_params_override = OmegaConf.to_container(scheduler_params_override, resolve=True)
# Get and instantiate the Config dataclass for this scheduler
scheduler_params_cls = get_scheduler_config(scheduler_params_name, **scheduler_params_override)
scheduler_params = scheduler_params_cls # instantiate the parameters object
# extract just the dictionary from the Config object
scheduler_args = vars(scheduler_params)
# Extract value to monitor in losses, if provided.
if "monitor" in scheduler_config:
monitor = scheduler_config.get("monitor")
else:
# Default to train loss
monitor = "loss"
# Store exact max_steps if it is provided
if "max_steps" in scheduler_config and scheduler_config["max_steps"] is not None:
max_steps = scheduler_config["max_steps"]
elif "t_max_epochs" in scheduler_config:
# Compute effective max_steps if t_max_epochs is provided
if train_dataloader is None:
logging.warning(
"As `t_max_epochs` is provided/computed, it is required to pass the train dataloader in order\n"
"to compute effective maximum number of steps.\n"
"Scheduler will not be instantiated !"
)
return None
# Raise exception if neither `max_steps` nor `t_max_epochs` is provided
if scheduler_config.get("t_max_epochs", None) is None:
logging.warning(
"`t_max_epochs` cannot be None when `max_steps` is not not provided.\n"
"This can occur when `train dataloader` is not available to correctly "
"prepare the scheduler.\n"
"Scheduler will not be instantiated !"
)
return None
# Get iters_per_batch
max_epochs = scheduler_config.get("t_max_epochs")
accumulate_grad_batches = scheduler_config.get("t_accumulate_grad_batches")
limit_train_batches = scheduler_config.get("t_limit_train_batches")
num_workers = scheduler_config.get("t_num_workers")
# Compute effective num max_steps
num_samples = len(train_dataloader.dataset)
# we may need to override ModelPT setup_optimization
if train_dataloader.batch_size is not None:
batch_size = train_dataloader.batch_size
elif hasattr(train_dataloader, "batch_sampler") and train_dataloader.batch_sampler is not None:
if train_dataloader.batch_sampler.micro_batch_size is not None:
batch_size = train_dataloader.batch_sampler.micro_batch_size
else:
raise ValueError(f"Could not find batch_size from batch_sampler: {train_dataloader.batch_sampler}")
else:
raise ValueError(f"Could not find batch_size from train_dataloader: {train_dataloader}")
drop_last = train_dataloader.drop_last
max_steps = compute_max_steps(
max_epochs=max_epochs,
accumulate_grad_batches=accumulate_grad_batches,
limit_train_batches=limit_train_batches,
num_workers=num_workers,
num_samples=num_samples,
batch_size=batch_size,
drop_last=drop_last,
)
else:
logging.warning(
"Neither `max_steps` nor `iters_per_batch` were provided to `optim.sched`, "
"cannot compute effective `max_steps` !\n"
"Scheduler will not be instantiated !"
)
return None
# Inject max_steps (effective or provided) into the scheduler config
scheduler_args["max_steps"] = max_steps
# Get the scheduler class from the config
scheduler_cls = get_scheduler(scheduler_name, **scheduler_args)
# Pop 'max_steps' if it's not required by the scheduler
if "max_steps" not in inspect.signature(scheduler_cls).parameters:
scheduler_args.pop("max_steps")
# Instantiate the LR schedule
schedule = scheduler_cls(optimizer, **scheduler_args)
logging.info(
'Scheduler "%s" \nwill be used during training (effective maximum steps = %d) - \nParameters : \n(%s)',
str(schedule),
max_steps,
OmegaConf.to_yaml(OmegaConf.create(scheduler_args)),
)
# Wrap the schedule in PTL arguments to perform stepwise computation. Rather than epoch level computation.
reduce_lr_on_plateau = bool(isinstance(schedule, optim.lr_scheduler.ReduceLROnPlateau))
schedule_dict = {
"scheduler": schedule,
"interval": interval,
"frequency": 1,
"monitor": monitor,
"reduce_on_plateau": reduce_lr_on_plateau,
}
return schedule_dict
def compute_max_steps(
max_epochs, accumulate_grad_batches, limit_train_batches, num_workers, num_samples, batch_size, drop_last
):
"""Compute effective max_steps from the provided parameters.
Parameters
----------
max_epochs : int
Maximum number of epochs to train for.
accumulate_grad_batches : int
Number of batches to accumulate gradients for.
limit_train_batches : int
Number of batches to train for.
num_workers : int
Number of workers to use for training.
num_samples : int
Number of samples in the dataset.
batch_size : int
Batch size.
drop_last : bool
Whether to drop the last batch or not.
"""
_round = math.floor if drop_last else math.ceil
sampler_num_samples = math.ceil(num_samples / max(1, num_workers))
if drop_last and num_workers > 1:
logging.warning(
"Please note that drop_last is broken in pytorch 1.6.0. We will fix when pytorch 1.7.0 is released"
)
# TODO: Master version, not in pytorch 1.6.0
steps_per_epoch = _round(sampler_num_samples / batch_size)
if isinstance(limit_train_batches, int) or limit_train_batches == 0.0:
steps_per_epoch = min(steps_per_epoch, int(limit_train_batches))
elif steps_per_epoch != float("inf"):
# limit_train_batches is a percentage of batches per epoch
steps_per_epoch = int(steps_per_epoch * limit_train_batches)
return math.ceil(steps_per_epoch / accumulate_grad_batches) * max_epochs
AVAILABLE_SCHEDULERS = {
"WarmupPolicy": WarmupPolicy,
"WarmupHoldPolicy": WarmupHoldPolicy,
"SquareAnnealing": SquareAnnealing,
"CosineAnnealing": CosineAnnealing,
"NoamAnnealing": NoamAnnealing,
"NoamHoldAnnealing": NoamHoldAnnealing,
"WarmupAnnealing": WarmupAnnealing,
"InverseSquareRootAnnealing": InverseSquareRootAnnealing,
"T5InverseSquareRootAnnealing": T5InverseSquareRootAnnealing,
"SquareRootAnnealing": SquareRootAnnealing,
"PolynomialDecayAnnealing": PolynomialDecayAnnealing,
"PolynomialHoldDecayAnnealing": PolynomialHoldDecayAnnealing,
"StepLR": pt_scheduler.StepLR,
"ExponentialLR": pt_scheduler.ExponentialLR,
"ReduceLROnPlateau": pt_scheduler.ReduceLROnPlateau,
"CyclicLR": pt_scheduler.CyclicLR,
}
EPOCH_SCHEDULERS = {
"ExponentialLR": pt_scheduler.ExponentialLR,
"ReduceLROnPlateau": pt_scheduler.ReduceLROnPlateau,
}
