# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Optax Schedules.
Schedules may be used to anneal the value of a hyper-parameter over time; for
instance, they may be used to anneal the learning rate used to update an agent's
parameters or the exploration factor used to select actions.
"""
from typing import Iterable, Optional, Union
from absl import logging
import chex
import jax
import jax.numpy as jnp
import numpy as np
from optax._src import base
from optax.schedules import _join
def constant_schedule(value: Union[float, int]) -> base.Schedule:
"""Constructs a constant schedule.
Args:
value: value to be held constant throughout.
Returns:
schedule
A function that maps step counts to values.
Examples:
>>> schedule_fn = optax.constant_schedule(5)
>>> schedule_fn(0)
5
>>> schedule_fn(100)
5
"""
return lambda count: value
def polynomial_schedule(
init_value: chex.Scalar,
end_value: chex.Scalar,
power: chex.Scalar,
transition_steps: int,
transition_begin: int = 0,
) -> base.Schedule:
r"""Constructs a schedule with polynomial transition from init to end value.
This function transitions the learning rate from an initial value
(``init_value``) to a final value (``end_value``) over a specified number of
steps (``transition_steps``) with a polynomial function of power ``power``.
The transition can optionally begin after a specified number of initial steps
(``transition_begin``).
More precisely, the learning rate at iteration :math:`t` is given by:
.. math::
\begin{cases}
I, & \text{if } t < B \\
(I - E) \left( 1 - \frac{t - B}{T} \right)^{P} + E, &
\text{if } B \leq t < B + T \\
E, & \text{if } t \geq B + T
\end{cases}
where :math:`I` is the initial value, :math:`E` is the end value,
:math:`B` is the transition begin, :math:`T` is the transition steps,
and :math:`P` is the power used for the polynomial transition.
Args:
init_value: initial value for the scalar to be annealed.
end_value: end value of the scalar to be annealed.
power: the power of the polynomial used to transition from init to end.
transition_steps: number of steps over which annealing takes place.
The scalar starts changing at ``transition_begin`` steps and completes
the transition by ``transition_begin + transition_steps`` steps.
If ``transition_steps <= 0``, then the entire annealing process is
disabled and the value is held fixed at ``init_value``.
transition_begin: must be positive. After how many steps to start annealing
(before this many steps the scalar value is held fixed at ``init_value``).
Returns:
schedule
A function that maps step counts to values.
Examples:
>>> schedule_fn = optax.polynomial_schedule(
... init_value=1.0, end_value=0.01, transition_steps=100, power=2)
>>> schedule_fn(0) # learning rate on the first iteration
Array(1., dtype=float32, weak_type=True)
>>> schedule_fn(100) # learning rate on the last iteration
Array(0.01, dtype=float32, weak_type=True)
The following example uses a non-zero ``transition_begin``. In this case
the learning rate is kept constant for the first ``transition_begin``
iterations:
>>> schedule_fn = optax.polynomial_schedule(
... init_value=1.0,
... end_value=0.01,
... transition_steps=100,
... transition_begin=5,
... power=2,
... )
>>> counts = [0, 5, 6, 104, 105, 110]
>>> print(
... *[f'count:{i} value:{schedule_fn(i):.4f}' for i in counts],
... sep='\n')
count:0 value:1.0000
count:5 value:1.0000
count:6 value:0.9803
count:104 value:0.0101
count:105 value:0.0100
count:110 value:0.0100
"""
if transition_steps <= 0:
logging.info(
'A polynomial schedule was set with a non-positive `transition_steps` '
'value; this results in a constant schedule with value `init_value`.'
)
return lambda count: init_value
if transition_begin < 0:
logging.info(
'A polynomial schedule was set with a negative `transition_begin` '
'value; this will result in `transition_begin` falling back to `0`.'
)
transition_begin = 0
def schedule(count):
count = jnp.clip(count - transition_begin, 0, transition_steps)
frac = 1 - count / transition_steps
return (init_value - end_value) * (frac**power) + end_value
return schedule
def linear_schedule(
init_value: chex.Scalar,
end_value: chex.Scalar,
transition_steps: int,
transition_begin: int = 0,
) -> base.Schedule:
r"""Schedule with linear transition from ``init_value`` to ``end_value``.
More precisely, the learning rate at iteration :math:`t` is given by:
.. math::
\begin{cases}
I, & \text{if } t < B \\
I + \frac{t - B}{T} (E - I), & \text{if } B \leq t < B + T \\
E, & \text{if } t \geq B + T
\end{cases}
where :math:`I` is the initial value, :math:`E` is the end value,
:math:`B` is the transition begin, and :math:`T` is the transition steps.
This schedule is equivalent to :func:`optax.polynomial_schedule` with
``power=1``.
Args:
init_value: initial value for the scalar to be annealed.
end_value: end value of the scalar to be annealed.
transition_steps: number of steps over which annealing takes place. The
scalar starts changing at ``transition_begin`` steps and completes the
transition by ``transition_begin + transition_steps`` steps. If
``transition_steps <= 0``, then the entire annealing process is disabled
and the value is held fixed at ``init_value``.
transition_begin: must be positive. After how many steps to start annealing
(before this many steps the scalar value is held fixed at ``init_value``).
Returns:
schedule
A function that maps step counts to values.
Examples:
>>> schedule_fn = optax.linear_schedule(
... init_value=1.0, end_value=0.01, transition_steps=100)
>>> schedule_fn(0) # learning rate on the first iteration
Array(1., dtype=float32, weak_type=True)
>>> schedule_fn(100) # learning rate on the last iteration
Array(0.01, dtype=float32, weak_type=True)
"""
return polynomial_schedule(
init_value=init_value,
end_value=end_value,
power=1,
transition_steps=transition_steps,
transition_begin=transition_begin,
)
def piecewise_constant_schedule(
init_value: float, boundaries_and_scales: Optional[dict[int, float]] = None
) -> base.Schedule:
"""Piecewise constant schedule with scaled jumps at specific boundaries.
At each step `t`, this schedule returns `init_value` scaled by the product
of all factors `f_i` such that `t >= b_i`, where `(b_i, f_i)` are the
entries in `boundaries_and_scales`.
Args:
init_value: The starting value of the schedule.
boundaries_and_scales: Dictionary of `{step: scale}` where `scale` is
multiplied into the schedule value at the given `step`. All `scale` values
must be non-negative.
Returns:
A function that maps step index to schedule value.
Example:
>>> sched = optax.piecewise_constant_schedule(
... init_value=1.0, boundaries_and_scales={100: 0.1, 200: 0.01})
>>> print(sched(50)) # before first boundary
1.0
>>> print(sched(150)) # after first boundary
0.1
>>> print(sched(250)) # after second boundary
0.001
"""
if boundaries_and_scales is not None:
all_positive = all(scale >= 0.0 for scale in boundaries_and_scales.values())
if not all_positive:
raise ValueError(
'`piecewise_constant_schedule` expects non-negative scale factors'
)
def schedule(count):
v = init_value
if boundaries_and_scales is not None:
for threshold, scale in sorted(boundaries_and_scales.items()):
indicator = jnp.maximum(0.0, jnp.sign(threshold - count))
v = v * indicator + (1 - indicator) * scale * v
return v
return schedule
def exponential_decay(
init_value: float,
transition_steps: int,
decay_rate: float,
transition_begin: int = 0,
staircase: bool = False,
end_value: Optional[float] = None,
) -> base.Schedule:
"""Constructs a schedule with either continuous or discrete exponential decay.
This function applies an exponential decay function to a provided initial
value. When ``count >= transition_begin`` the function returns the decayed
value as:
.. code-block::
rate_factor = ((count - transition_begin) / transition_steps)
decayed_value = init_value * (decay_rate ** rate_factor)
If the argument ``staircase`` is ``True`` then ``count / transition_steps`` is
an integer division and the decayed value follows a staircase function.
Args:
init_value: the initial learning rate.
transition_steps: must be positive. See the decay computation above.
decay_rate: must not be zero. The decay rate.
transition_begin: must be positive. After how many steps to start annealing
(before this many steps the scalar value is held fixed at `init_value`).
staircase: if ``True``, decay the values at discrete intervals.
end_value: the value at which the exponential decay stops. When ``decay_rate
< 1``, ``end_value`` is treated as a lower bound, otherwise as an upper
bound. Has no effect when ``decay_rate = 0``.
Returns:
schedule
A function that maps step counts to values.
"""
if transition_steps <= 0:
logging.info(
'An exponential schedule was set with a non-positive `transition_steps`'
' value; this will result in a constant schedule with value '
'`init_value`.'
)
return lambda count: init_value
if decay_rate == 0:
logging.info(
'An exponential schedule was set with a zero `decay_rate` value; '
'this will result in a constant schedule with value `init_value`.'
)
return lambda count: init_value
if transition_begin < 0:
logging.info(
'An exponential schedule was set with a negative `transition_begin` '
'value; this will result in `transition_begin` falling back to `0`.'
)
transition_begin = 0
if end_value is not None:
clip_fn = jnp.maximum if decay_rate < 1.0 else jnp.minimum
def schedule(count):
decreased_count = count - transition_begin
p = decreased_count / transition_steps
if staircase:
p = jnp.floor(p)
decayed_value = jnp.where(
decreased_count <= 0, init_value, init_value * jnp.power(decay_rate, p)
)
if end_value is not None:
decayed_value = clip_fn(decayed_value, end_value) # pylint: disable=undefined-variable
return decayed_value
return schedule
def cosine_decay_schedule(
init_value: float,
decay_steps: int,
alpha: float = 0.0,
exponent: float = 1.0,
) -> base.Schedule:
r"""Returns a function which implements cosine learning rate decay.
This schedule smoothly decreases the learning rate over a specified number of
steps (``decay_steps``). The decay follows a cosine function, with an optional
exponent to modify the decay curve. A minimum value (``alpha``) ensures the
learning rate does not drop entirely to zero.
More precisely, the learning rate at iteration :math:`t` is given by:
.. math::
\begin{cases}
\alpha I + (1 - \alpha) I \left[ \frac{1}{2} \left( 1+ \cos \left( \pi\,\frac{t}{T} \right) \right) \right] ^p & \text{, if } t \leq T \\
\alpha I & \text{, if } t > T
\end{cases}
where :math:`T` is the number of decay steps (``decay_steps``), :math:`p` is
the ``exponent`` and :math:`I` is the initial value (``init_value``).
Args:
init_value: An initial value for the learning rate.
decay_steps: Positive integer - the number of steps for which to apply
the decay for.
alpha: The minimum value of the multiplier used to adjust the
learning rate. Defaults to 0.0.
exponent: The default decay is ``0.5 * (1 + cos(pi * t/T))``, where
``t`` is the current timestep and ``T`` is the ``decay_steps``. The
exponent modifies this to be ``(0.5 * (1 + cos(pi * t/T))) ** exponent``.
Defaults to 1.0.
Returns:
schedule
A function that maps step counts to values.
References:
Loshchilov et al., `SGDR: Stochastic Gradient Descent with Warm Restarts
<https://arxiv.org/abs/1608.03983>`_, 2017
""" # noqa: E501
if not decay_steps > 0:
raise ValueError(
'The cosine_decay_schedule requires positive decay_steps, got'
f' {decay_steps=}.'
)
def schedule(count):
# Avoid int -> int32 overflow in jitted code.
nonlocal decay_steps
decay_steps, count = jax.tree.map(
lambda x: float(x) if isinstance(x, int) else x, (decay_steps, count)
)
count = jnp.minimum(count, decay_steps)
cosine_decay = 0.5 * (1 + jnp.cos(jnp.pi * count / decay_steps))
decayed = (1 - alpha) * cosine_decay**exponent + alpha
return init_value * decayed
return schedule
def _linear_interpolate(start: float, end: float, pct: float):
"""Linearly interpolate between two values."""
return (end - start) * pct + start
def _cosine_interpolate(start: float, end: float, pct: float):
"""Cosine interpolate between two values (smoother transitions)."""
return end + (start - end) / 2.0 * (jnp.cos(jnp.pi * pct) + 1)
def piecewise_interpolate_schedule(
interpolate_type: str,
init_value: float,
boundaries_and_scales: Optional[dict[int, float]] = None,
) -> base.Schedule:
"""Piecewise interpolated schedule with linear or cosine transitions.
This schedule interpolates smoothly between values scaled at specified
step boundaries. The interpolation occurs **between the accumulated scaled
values**, not the raw multiplicative scales themselves.
At each boundary, the scaling factor is multiplied into the current value.
Between these boundaries, the schedule applies either linear or cosine
interpolation.
Args:
interpolate_type: Either `'linear'` or `'cosine'`, specifying the
interpolation method used between boundary segments.
init_value: Starting value for the schedule.
boundaries_and_scales: A dictionary `{step: scale}` that defines the
boundaries and their corresponding multiplicative scaling factors.
Returns:
A function that maps step counts to interpolated values.
Example:
>>> sched = optax.piecewise_interpolate_schedule(
... interpolate_type='linear',
... init_value=1.0,
... boundaries_and_scales={4: 0.1, 8: 0.01}
... )
>>> for step in range(0, 11, 2):
... print(f"Step {step}: {sched(step):.6f}")
Step 0: 1.000000
Step 2: 0.550000
Step 4: 0.100000
Step 6: 0.050500
Step 8: 0.001000
Step 10: 0.001000
.. note::
This schedule accumulates scaling factors and then interpolates between
those accumulated values. It does *not* interpolate directly between the
provided scales. This behavior can appear counterintuitive but allows
for smooth and controlled transitions in learning rates.
"""
if interpolate_type == 'linear':
interpolate_fn = _linear_interpolate
elif interpolate_type == 'cosine':
interpolate_fn = _cosine_interpolate
else:
raise ValueError("`interpolate_type` must be either 'cosine' or 'linear'")
if boundaries_and_scales:
boundaries, scales = zip(*sorted(boundaries_and_scales.items()))
if not all(scale >= 0.0 for scale in scales):
raise ValueError(
'`piecewise_interpolate_schedule` expects non-negative scale factors'
)
else:
boundaries, scales = (), ()
bounds = np.stack((0,) + boundaries)
values = np.cumprod(np.stack((init_value,) + scales))
interval_sizes = bounds[1:] - bounds[:-1]
def schedule(count):
indicator = (bounds[:-1] <= count) & (count < bounds[1:])
pct = (count - bounds[:-1]) / interval_sizes
interp_vals = interpolate_fn(values[:-1], values[1:], pct)
return indicator.dot(interp_vals) + (bounds[-1] <= count) * values[-1]
return schedule
def linear_onecycle_schedule(
transition_steps: int,
peak_value: float,
pct_start: float = 0.3,
pct_final: float = 0.85,
div_factor: float = 25.0,
final_div_factor: float = 1e4,
) -> base.Schedule:
r"""Returns a learning rate with three linear phases.
* *Phase 1*, from iteration 0 to ``pct_start * transition_steps``. The
learning rate increases linearly from ``peak_value / div_factor`` to
``peak_value``.
* *Phase 2*, from iteration ``pct_start * transition_steps`` to
``pct_final * transition_steps``. The learning rate decreases linearly from
``peak_value`` back to the initial ``peak_value/div_factor``.
* *Phase 3*: For the remaining steps, the learning rate interpolates between
``peak_value/div_factor`` and ``peak_value / final_div_factor``. If
``final_div_factor`` is larger than ``div_factor``, this is a decreasing
phase.
Args:
transition_steps: Number of steps over which annealing takes place.
peak_value: Maximum value attained by schedule at pct_start percent of the
cycle (in number of steps).
pct_start: The percentage of the cycle (in number of steps) spent increasing
the learning rate.
pct_final: The percentage of the cycle (in number of steps) spent increasing
to ``peak_value`` then decreasing back to ``init_value``.
div_factor: Determines the initial value via ``init_value = peak_value /
div_factor``.
final_div_factor: Determines the final value via ``final_value = init_value
/ final_div_factor``.
Returns:
schedule
A function that maps step counts to values
References:
Smith et al, `Super-Convergence: Very Fast Training of Neural Networks Using
Large Learning Rates <https://arxiv.org/abs/1708.07120>`_, 2017
"""
if transition_steps <= 0:
raise ValueError(
'A linear onecycle schedule was set with a non-positive '
'`transition_steps`'
)
return piecewise_interpolate_schedule(
'linear',
peak_value / div_factor,
{
int(pct_start * transition_steps): div_factor,
int(pct_final * transition_steps): 1.0 / div_factor,
transition_steps: 1.0 / final_div_factor,
},
)
def cosine_onecycle_schedule(
transition_steps: int,
peak_value: float,
pct_start: float = 0.3,
div_factor: float = 25.0,
final_div_factor: float = 1e4,
) -> base.Schedule:
"""Returns a function which implements the onecycle learning rate schedule.
This schedule increases the learning rate and then decreases it in a
cosine-like manner. The number of steps over which the learning rate increases
is determined by the ``pct_start`` argument. The maximum value of the learning
rate is determined by the ``peak_value`` argument, the initial value of the
learning rate is determined through the formula ``init_value = peak_value /
div_factor``, and the final value is determined by the ``final_div_factor``
argument.
Args:
transition_steps: Number of steps over which annealing takes place.
peak_value: Maximum value attained by schedule at pct_start percent of the
cycle (in number of steps).
pct_start: The percentage of the cycle (in number of steps) spent increasing
the learning rate.
div_factor: Determines the initial value via ``init_value = peak_value /
div_factor``.
final_div_factor: Determines the final value via ``final_value = init_value
/ final_div_factor``.
Returns:
schedule
A function that maps step counts to values
References:
Smith et al, `Super-Convergence: Very Fast Training of Neural Networks Using
Large Learning Rates <https://arxiv.org/abs/1708.07120>`_, 2017
"""
if transition_steps <= 0:
raise ValueError(
'A linear onecycle schedule was set with a non-positive '
'`transition_steps`'
)
return piecewise_interpolate_schedule(
'cosine',
peak_value / div_factor,
{
int(pct_start * transition_steps): div_factor,
int(transition_steps): 1.0 / (div_factor * final_div_factor),
},
)
def warmup_constant_schedule(
init_value: float,
peak_value: float,
warmup_steps: int,
) -> base.Schedule:
r"""Linear warmup followed by constant schedule i.e no decay.
Args:
init_value: Initial value for the scalar to be annealed.
peak_value: Peak value for scalar to be annealed at end of warmup.
warmup_steps: Positive integer, the length of the linear warmup.
Returns:
schedule
A function that maps step counts to values
"""
return linear_schedule(
init_value=init_value,
end_value=peak_value,
transition_steps=warmup_steps,
)
def warmup_cosine_decay_schedule(
init_value: float,
peak_value: float,
warmup_steps: int,
decay_steps: int,
end_value: float = 0.0,
exponent: float = 1.0,
) -> base.Schedule:
r"""Linear warmup followed by cosine decay.
Args:
init_value: Initial value for the scalar to be annealed.
peak_value: Peak value for scalar to be annealed at end of warmup.
warmup_steps: Positive integer, the length of the linear warmup.
decay_steps: Positive integer, the total length of the schedule. Note that
this includes the warmup time, so the number of steps during which cosine
annealing is applied is ``decay_steps - warmup_steps``.
end_value: End value of the scalar to be annealed.
exponent: The default decay is ``0.5 * (1 + cos(pi t/T))``, where ``t`` is
the current timestep and ``T`` is ``decay_steps``. The exponent modifies
this to be ``(0.5 * (1 + cos(pi * t/T))) ** exponent``. Defaults to 1.0.
Returns:
schedule
A function that maps step counts to values
"""
alpha = 0.0 if peak_value == 0.0 else end_value / peak_value
schedules = [
linear_schedule(
init_value=init_value,
end_value=peak_value,
transition_steps=warmup_steps,
),
cosine_decay_schedule(
init_value=peak_value,
decay_steps=decay_steps - warmup_steps,
alpha=alpha,
exponent=exponent,
),
]
return _join.join_schedules(schedules, [warmup_steps])
def warmup_exponential_decay_schedule(
init_value: float,
peak_value: float,
warmup_steps: int,
transition_steps: int,
decay_rate: float,
transition_begin: int = 0,
staircase: bool = False,
end_value: Optional[float] = None,
) -> base.Schedule:
"""Linear warmup followed by exponential decay.
Args:
init_value: Initial value for the scalar to be annealed.
peak_value: Peak value for scalar to be annealed at end of warmup.
warmup_steps: Positive integer, the length of the linear warmup.
transition_steps: must be positive. See :func:`optax.exponential_decay` for
more details.
decay_rate: must not be zero. The decay rate.
transition_begin: must be positive. After how many steps to start annealing
(before this many steps the scalar value is held fixed at ``peak_value``).
staircase: if ``True``, decay the values at discrete intervals.
end_value: the value at which the exponential decay stops. When ``decay_rate
< 1``, ``end_value`` is treated as a lower bound, otherwise as an upper
bound. Has no effect when ``decay_rate = 0``.
Returns:
schedule
A function that maps step counts to values
"""
schedules = [
linear_schedule(
init_value=init_value,
end_value=peak_value,
transition_steps=warmup_steps,
),
exponential_decay(
init_value=peak_value,
transition_steps=transition_steps,
decay_rate=decay_rate,
transition_begin=transition_begin,
staircase=staircase,
end_value=end_value,
),
]
return _join.join_schedules(schedules, [warmup_steps])
def sgdr_schedule(
cosine_kwargs: Iterable[dict[str, chex.Numeric]],
) -> base.Schedule:
"""SGD with warm restarts.
This learning rate schedule applies multiple joined cosine decay cycles.
Args:
cosine_kwargs: An Iterable of dicts, where each element specifies the
arguments to pass to each cosine decay cycle. The ``decay_steps`` kwarg
will specify how long each cycle lasts for, and therefore when to
transition to the next cycle.
Returns:
schedule
A function that maps step counts to values
References:
Loshchilov et al., `SGDR: Stochastic Gradient Descent with Warm Restarts
<https://arxiv.org/abs/1608.03983>`_, 2017
"""
boundaries = []
schedules = []
step = 0
for kwargs in cosine_kwargs:
schedules += [warmup_cosine_decay_schedule(**kwargs)]
boundaries += [step + kwargs['decay_steps']]
step += kwargs['decay_steps']
return _join.join_schedules(schedules, boundaries[:-1])
