Repository: larq/larq

Issue:
Generic hyperparameter scheduler 
The Idea is to have a TF scheduler which is optimizer/hyperparameter agnostic. 
As an example, it can be used for gamma, threshold or learning rate of the floating point optimizer in Bop.


Discussion:

Resolving patch:
diff --git a/larq/callbacks.py b/larq/callbacks.py
--- a/larq/callbacks.py
+++ b/larq/callbacks.py
@@ -74,3 +74,45 @@ def on_train_begin(self, logs=None):
 
     def on_epoch_end(self, epoch, logs=None):
         self._maybe_log_and_store(self.epoch_previous_weights, logs)
+
+
+class HyperparameterScheduler(tf.keras.callbacks.Callback):
+    """Generic hyperparameter scheduler.
+
+    # Arguments:
+    schedule: a function that takes an epoch index as input
+        (integer, indexed from 0) and returns a new hyperparameter as output.
+    hyperparameter: str. the name of the hyperparameter to be scheduled.
+    verbose: int. 0: quiet, 1: update messages.
+    """
+
+    def __init__(self, schedule, hyperparameter, verbose=0):
+        super(HyperparameterScheduler, self).__init__()
+        self.schedule = schedule
+        self.hyperparameter = hyperparameter
+        self.verbose = verbose
+
+    def on_epoch_begin(self, epoch, logs=None):
+        if not hasattr(self.model.optimizer, self.hyperparameter):
+            raise ValueError(
+                f'Optimizer must have a "{self.hyperparameter}" attribute.'
+            )
+
+        hp = getattr(self.model.optimizer, self.hyperparameter)
+        try:  # new API
+            hyperparameter_val = tf.keras.backend.get_value(hp)
+            hyperparameter_val = self.schedule(epoch, hyperparameter_val)
+        except TypeError:  # Support for old API for backward compatibility
+            hyperparameter_val = self.schedule(epoch)
+
+        tf.keras.backend.set_value(hp, hyperparameter_val)
+
+        if self.verbose > 0:
+            print(
+                f"Epoch {epoch + 1}: {self.hyperparameter} changning to {tf.keras.backend.get_value(hp)}."
+            )
+
+    def on_epoch_end(self, epoch, logs=None):
+        logs = logs or {}
+        hp = getattr(self.model.optimizer, self.hyperparameter)
+        logs[self.hyperparameter] = tf.keras.backend.get_value(hp)
diff --git a/larq/optimizers_v1.py b/larq/optimizers_v1.py
--- a/larq/optimizers_v1.py
+++ b/larq/optimizers_v1.py
@@ -137,8 +137,8 @@ def __init__(self, fp_optimizer, threshold=1e-5, gamma=1e-2, name="Bop", **kwarg
         with tf.keras.backend.name_scope(self.__class__.__name__):
 
             self.fp_optimizer = fp_optimizer
-            self.threshold = threshold
-            self.gamma = gamma
+            self.threshold = tf.keras.backend.variable(threshold, name="threshold")
+            self.gamma = tf.keras.backend.variable(gamma, name="gamma")
 
     def get_updates(self, loss, params):
         grads = self.get_gradients(loss, params)
@@ -181,8 +181,8 @@ def __getattr__(self, name):
     def get_config(self):
         fp_optimizer_config = self.fp_optimizer.get_config()
         config = {
-            "threshold": self.threshold,
-            "gamma": self.gamma,
+            "threshold": float(tf.keras.backend.get_value(self.threshold)),
+            "gamma": float(tf.keras.backend.get_value(self.gamma)),
             "fp_optimizer": {
                 "class_name": self.fp_optimizer.__class__.__name__,
                 "config": fp_optimizer_config,


Test patch:
diff --git a/larq/callbacks_test.py b/larq/callbacks_test.py
--- a/larq/callbacks_test.py
+++ b/larq/callbacks_test.py
@@ -1,6 +1,12 @@
 import numpy as np
 import tensorflow as tf
 import larq as lq
+
+from larq import testing_utils as lq_testing_utils
+from larq.callbacks import HyperparameterScheduler
+
+from tensorflow import keras
+from tensorflow.python.keras import testing_utils
 from tensorflow.python.keras import keras_parameterized
 
 
@@ -46,3 +52,64 @@ def test_quantization_logger(self):
         model.fit(x, y, batch_size=1, epochs=3, callbacks=[logger, history])
         assert history.batches == [4, 9, 14, 19] * 3
         assert history.epochs == [0, 1, 2]
+
+
+class TestHyperparameterScheduler:
+    def test_hyper_parameter_scheduler(self):
+        np.random.seed(1337)
+        (x_train, y_train), (x_test, y_test) = testing_utils.get_test_data(
+            train_samples=1000, test_samples=0, input_shape=(10,), num_classes=2
+        )
+
+        y_train = keras.utils.to_categorical(y_train)
+
+        model = lq_testing_utils.get_small_bnn_model(
+            x_train.shape[1], 20, y_train.shape[1]
+        )
+        bop_optimizer = lq.optimizers.Bop(fp_optimizer=tf.keras.optimizers.Adam(0.01))
+        model.compile(
+            loss="categorical_crossentropy",
+            optimizer=bop_optimizer,
+            metrics=["accuracy"],
+        )
+
+        def scheduler(x):
+            return 1.0 / (1.0 + x)
+
+        cbk_gamma_scheduler = HyperparameterScheduler(
+            scheduler, hyperparameter="gamma", verbose=1
+        )
+        cbk_threshold_scheduler = HyperparameterScheduler(
+            scheduler, hyperparameter="threshold", verbose=1
+        )
+        cbk_lr_scheduler = HyperparameterScheduler(
+            scheduler, hyperparameter="lr", verbose=1
+        )
+
+        num_epochs = 10
+        model.fit(
+            x_train,
+            y_train,
+            epochs=num_epochs,
+            batch_size=16,
+            callbacks=[cbk_gamma_scheduler, cbk_lr_scheduler, cbk_threshold_scheduler],
+            verbose=0,
+        )
+
+        np.testing.assert_almost_equal(
+            keras.backend.get_value(model.optimizer.gamma),
+            scheduler(num_epochs - 1),
+            decimal=8,
+        )
+
+        np.testing.assert_almost_equal(
+            keras.backend.get_value(model.optimizer.threshold),
+            scheduler(num_epochs - 1),
+            decimal=8,
+        )
+
+        np.testing.assert_almost_equal(
+            keras.backend.get_value(model.optimizer.lr),
+            scheduler(num_epochs - 1),
+            decimal=8,
+        )