Added ReduceLROnPlateau callback for VI.

pymc/variational/callbacks.py

@@ -18,7 +18,7 @@
 
 import numpy as np
 
-__all__ = ["Callback", "CheckParametersConvergence", "Tracker"]
+__all__ = ["Callback", "CheckParametersConvergence", "ReduceLROnPlateau", "Tracker"]
 
 
 class Callback:
@@ -93,6 +93,77 @@
         return np.concatenate([sh.get_value().flatten() for sh in shared_list])
 
 
+class ReduceLROnPlateau(Callback):
+    """Reduce learning rate when the loss has stopped improving.
+
+    This is inspired by Keras' homonymous callback:
+    https://github.com/keras-team/keras/blob/v2.14.0/keras/callbacks.py
+
+    Parameters
+    ----------
+    optimizer: callable
+        PyMC optimizer
+    factor: float
+        factor by which the learning rate will be reduced: `new_lr = lr * factor`
+    patience: int
+        number of epochs with no improvement after which learning rate will be reduced
+    min_lr: float
+        lower bound on the learning rate
+    cooldown: int
+        number of iterations to wait before resuming normal operation after lr has been reduced
+    verbose: bool
+        false: quiet, true: update messages
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        factor=0.1,
+        patience=10,
+        min_lr=1e-6,
+        cooldown=0,
+    ):
+        self.optimizer = optimizer
+        self.factor = factor
+        self.patience = patience
+        self.min_lr = min_lr
+        self.cooldown = cooldown
+
+        self.cooldown_counter = 0
+        self.wait = 0
+        self.best = float("inf")
+
+    def __call__(self, approx, loss_hist, i):
+        current = loss_hist[-1]
+
+        if np.isinf(current):
+            return
+
+        if self.in_cooldown():
+            self.cooldown_counter -= 1
+            self.wait = 0
+            return
+
+        if current < self.best:
+            self.best = current
+            self.wait = 0
+        elif not np.isinf(self.best):
+            self.wait += 1
+            if self.wait >= self.patience:
+                self.reduce_lr()
+                self.cooldown_counter = self.cooldown
+                self.wait = 0
+
+    def reduce_lr(self):
+        old_lr = float(self.optimizer.keywords["learning_rate"])
+        if old_lr > self.min_lr:
+            new_lr = max(old_lr * self.factor, self.min_lr)
+            self.optimizer.keywords["learning_rate"] = new_lr
+
+    def in_cooldown(self):
+        return self.cooldown_counter > 0
+
+
 class Tracker(Callback):
     """
     Helper class to record arbitrary stats during VI

tests/variational/test_callbacks.py

@@ -52,3 +52,33 @@ def test_tracker_callback():
     tracker = pm.callbacks.Tracker(bad=lambda t: t)  # bad signature
     with pytest.raises(TypeError):
         tracker(None, None, 1)
+
+
+def test_reducelronplateau_callback():
+    optimizer = pm.adam(learning_rate=0.1)
+    cb = pm.variational.callbacks.ReduceLROnPlateau(
+        optimizer=optimizer,
+        patience=1,
+        min_lr=0.001,
+    )
+    cb(None, [float("inf")], 1)
+    np.testing.assert_almost_equal(optimizer.keywords["learning_rate"], 0.1)
+    assert cb.best == float("inf")
+    cb(None, [float("inf"), 2], 1)
+    np.testing.assert_almost_equal(optimizer.keywords["learning_rate"], 0.1)
+    assert cb.best == 2
+    cb(None, [float("inf"), 2, 1], 1)
+    np.testing.assert_almost_equal(optimizer.keywords["learning_rate"], 0.1)
+    assert cb.best == 1
+    cb(None, [float("inf"), 2, 1, 99], 1)
+    np.testing.assert_almost_equal(optimizer.keywords["learning_rate"], 0.01)
+    assert cb.best == 1
+    cb(None, [float("inf"), 2, 1, 99, 0.9], 1)
+    np.testing.assert_almost_equal(optimizer.keywords["learning_rate"], 0.01)
+    assert cb.best == 0.9
+    cb(None, [float("inf"), 2, 1, 99, 0.9, 99], 1)
+    np.testing.assert_almost_equal(optimizer.keywords["learning_rate"], 0.001)
+    assert cb.best == 0.9
+    cb(None, [float("inf"), 2, 1, 99, 0.9, 99, 99], 1)
+    np.testing.assert_almost_equal(optimizer.keywords["learning_rate"], 0.001)
+    assert cb.best == 0.9