Repository: google-deepmind/optax

Issue:
Implement OneCycle Learning Rate Schedule
OneCycle learning rate from this paper: https://arxiv.org/abs/1708.07120. 

The paper describes a 3-stage piecewise linear schedule. [PyTorch](https://pytorch.org/docs/stable/optim.html#torch.optim.lr_scheduler.OneCycleLR) and [FastAI](https://fastai1.fast.ai/callbacks.one_cycle.html) implement a 2 stage variant with cosine annealing (which apparently works better in practice).

I would like to contribute both to this library, if you think they are suitable additions.


Discussion:
This would be a beautiful contribution thanks!

Resolving patch:
diff --git a/optax/__init__.py b/optax/__init__.py
--- a/optax/__init__.py
+++ b/optax/__init__.py
@@ -30,8 +30,11 @@
 from optax._src.control_variates import moving_avg_baseline
 from optax._src.schedule import constant_schedule
 from optax._src.schedule import cosine_decay_schedule
+from optax._src.schedule import cosine_onecycle_schedule
 from optax._src.schedule import exponential_decay
+from optax._src.schedule import linear_onecycle_schedule
 from optax._src.schedule import piecewise_constant_schedule
+from optax._src.schedule import piecewise_interpolate_schedule
 from optax._src.schedule import polynomial_schedule
 from optax._src.second_order import fisher_diag
 from optax._src.second_order import hessian_diag
@@ -109,6 +112,7 @@
     "control_delta_method",
     "control_variates_jacobians",
     "cosine_decay_schedule",
+    "cosine_onecycle_schedule",
     "exponential_decay",
     "fisher_diag",
     "flatten",
@@ -119,6 +123,7 @@
     "identity",
     "InitUpdate",
     "lamb",
+    "linear_onecycle_schedule",
     "measure_valued_jacobians",
     "moving_avg_baseline",
     "multi_normal",
@@ -127,6 +132,7 @@
     "Params",
     "pathwise_jacobians",
     "piecewise_constant_schedule",
+    "piecewise_interpolate_schedule",
     "polynomial_schedule",
     "rmsprop",
     "scale",
diff --git a/optax/_src/schedule.py b/optax/_src/schedule.py
--- a/optax/_src/schedule.py
+++ b/optax/_src/schedule.py
@@ -20,7 +20,7 @@
 parameters or the exploration factor used to select actions.
 """
 
-from typing import Dict, Union
+from typing import Dict, Union, Optional
 
 from absl import logging
 import jax.numpy as jnp
@@ -84,7 +84,7 @@ def schedule(count):
 
 def piecewise_constant_schedule(
     init_value: float,
-    boundaries_and_scales: Dict[int, float] = None):
+    boundaries_and_scales: Optional[Dict[int, float]] = None):
   """Returns a function which implements a piecewise constant schedule.
 
   Args:
@@ -203,3 +203,138 @@ def schedule(count):
     return init_value * decayed
 
   return schedule
+
+
+def _linear_interpolate(start: float, end: float, pct: float):
+  return (end-start) * pct + start
+
+
+def _cosine_interpolate(start: float, end: float, pct: float):
+  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):
+  """Returns a function which implements a piecewise interpolated schedule.
+
+  Args:
+    interpolate_type: 'linear' or 'cosine', specifying the interpolation
+      strategy.
+    init_value: An initial value `init_v`.
+    boundaries_and_scales: A map from boundaries `b_i` to non-negative scaling
+      factors `f_i`. At boundary step `b_i`, the schedule returns `init_v`
+      scaled by the product of all factors `f_j` such that `b_j` < `b_i`. The
+      values in between each boundary will be interpolated as per `type`.
+
+  Returns:
+    schedule: A function that maps step counts to values.
+  """
+  if interpolate_type == 'linear':
+    interpolate_fn = _linear_interpolate
+  elif interpolate_type == 'cosine':
+    interpolate_fn = _cosine_interpolate
+  else:
+    raise ValueError(
+        'The `piecewise_interpolate_schedule` expects \'cos\' or \'linear\' '
+        'for `interpolate_type`')
+
+  if boundaries_and_scales:
+    boundaries, scales = zip(*sorted(boundaries_and_scales.items()))
+    if not all(scale >= 0. for scale in scales):
+      raise ValueError(
+          'The `piecewise_interpolate_schedule` expects non-negative scale '
+          'factors')
+  else:
+    boundaries, scales = (), ()
+
+  bounds = jnp.stack((0,) + boundaries)
+  values = jnp.cumprod(jnp.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):
+  """Returns a function which implements the onecycle learning rate schedule.
+
+  This function uses a linear annealing strategy.
+  For more details see: https://arxiv.org/abs/1708.07120
+
+  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.
+  """
+  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. / div_factor,
+       transition_steps: 1. / 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):
+  """Returns a function which implements the onecycle learning rate schedule.
+
+  This function uses a cosine annealing strategy.
+  For more details see: https://arxiv.org/abs/1708.07120
+
+  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.
+  """
+  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. / (div_factor * final_div_factor)})


Test patch:
diff --git a/optax/_src/schedule_test.py b/optax/_src/schedule_test.py
--- a/optax/_src/schedule_test.py
+++ b/optax/_src/schedule_test.py
@@ -302,5 +302,90 @@ def test_decay_count_greater_count_with_alpha(self):
         np.array(generated_vals), atol=1e-3)
 
 
+class PiecewiseInterpolateTest(chex.TestCase):
+
+  @chex.all_variants()
+  def test_linear_piecewise(self):
+    schedule_fn = self.variant(schedule.piecewise_interpolate_schedule(
+        'linear', 200., {5: 1.5, 10: 0.25}))
+    generated_vals = [schedule_fn(step) for step in range(13)]
+    expected_vals = [200., 220., 240., 260., 280., 300., 255., 210., 165.,
+                     120., 75., 75., 75.]
+    np.testing.assert_allclose(generated_vals, expected_vals, atol=1e-3)
+
+  @chex.all_variants()
+  def test_cos_piecewise(self):
+    schedule_fn = self.variant(schedule.piecewise_interpolate_schedule(
+        'cosine', 400., {5: 1.2, 3: 0.6, 7: 1.}))
+    generated_vals = [schedule_fn(step) for step in range(9)]
+    expected_vals = [400., 360., 280., 240., 264., 288., 288., 288., 288.]
+    np.testing.assert_allclose(generated_vals, expected_vals, atol=1e-3)
+
+  @chex.all_variants()
+  def test_empty_dict(self):
+    schedule_fn = self.variant(schedule.piecewise_interpolate_schedule(
+        'linear', 13., {}))
+    generated_vals = [schedule_fn(step) for step in range(5)]
+    expected_vals = [13., 13., 13., 13., 13.]
+    np.testing.assert_allclose(generated_vals, expected_vals, atol=1e-3)
+
+  @chex.all_variants()
+  def test_no_dict(self):
+    schedule_fn = self.variant(schedule.piecewise_interpolate_schedule(
+        'cosine', 17.))
+    generated_vals = [schedule_fn(step) for step in range(3)]
+    expected_vals = [17., 17., 17.]
+    np.testing.assert_allclose(generated_vals, expected_vals, atol=1e-3)
+
+  def test_invalid_type(self):
+    with self.assertRaises(ValueError):
+      schedule.piecewise_interpolate_schedule('linar', 13.)
+    with self.assertRaises(ValueError):
+      schedule.piecewise_interpolate_schedule('', 13., {5: 3.})
+    with self.assertRaises(ValueError):
+      schedule.piecewise_interpolate_schedule(None, 13., {})
+
+  def test_invalid_scale(self):
+    with self.assertRaises(ValueError):
+      schedule.piecewise_interpolate_schedule('linear', 13., {5: -3})
+
+
+class OneCycleTest(chex.TestCase):
+
+  @chex.all_variants()
+  def test_linear(self):
+    schedule_fn = self.variant(schedule.linear_onecycle_schedule(
+        transition_steps=10,
+        peak_value=1000,
+        pct_start=0.3,
+        pct_final=0.7,
+        div_factor=10.,
+        final_div_factor=100.))
+
+    generated_vals = [schedule_fn(step) for step in range(12)]
+    expected_vals = [100., 400., 700., 1000., 775., 550., 325., 100., 67.,
+                     34., 1., 1.]
+    np.testing.assert_allclose(generated_vals, expected_vals, atol=1e-3)
+
+  @chex.all_variants()
+  def test_cosine(self):
+    schedule_fn = self.variant(schedule.cosine_onecycle_schedule(
+        transition_steps=5,
+        peak_value=1000.,
+        pct_start=0.4,
+        div_factor=10.,
+        final_div_factor=100.))
+
+    generated_vals = [schedule_fn(step) for step in range(7)]
+    expected_vals = [100., 550., 1000., 750.25, 250.75, 1., 1.]
+    np.testing.assert_allclose(generated_vals, expected_vals, atol=1e-3)
+
+  def test_nonpositive_transition_steps(self):
+    with self.assertRaises(ValueError):
+      schedule.cosine_onecycle_schedule(transition_steps=0, peak_value=5.)
+    with self.assertRaises(ValueError):
+      schedule.linear_onecycle_schedule(transition_steps=0, peak_value=5.)
+
+
 if __name__ == '__main__':
   absltest.main()