Repository: larq/larq

Issue:
Make the HyperparameterScheduler compatible with the CaseOptimizer
### Feature motivation
The HyperparameterScheduler is not compatible with the CaseOptimizer since the hyperparameters are attributes of the optimizers inside the CaseOptimizer. 

### Feature description
I propose one of the two possible solutions: 
Either we could give HyperparameterScheduler the optimizer as an argument. It could be called via ```  HyperparameterScheduler(schedule, hyperparameter, optimizer, verbose=0) ``` and the right optimizer inside the CaseOptimizer can be addressed. (@koenhelwegen) 

My second proposal would be to search the CaseOptimizer for optimizers that have the hyperparameter as attribute. Then the schedule can be applied to this optimizer only. The downside of this would be that in case there are two optimizers inside the CaseOptimizer that have a hyperparameter with the same name the schedule would be applied to both of them. I do not think this would happen very often but it could definitively be an issue. See code below for my second proposal. 
### Feature implementation
``` python 
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):
        for op in self.model.optimizer.optimizers:
            if hasattr(op, self.hyperparameter):

                hp = getattr(op, 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 {}
        for op in self.model.optimizer.optimizers:
            if hasattr(op, self.hyperparameter):
                hp = getattr(op, self.hyperparameter)
                logs[self.hyperparameter] = tf.keras.backend.get_value(hp)
```


Discussion:
I'm leaning towards the first option because it's a bit more explicit in what's happening; but one drawback could be that I'm not sure if/how it'd work when the optimizers get reloaded on an experiment restart/resume. So I'm also good with doing option two if we document it well.
Resuming and Restarting works fine with the first option. I am leaning to the first option as well since it gives more flexibility. I will make a PR. 

Resolving patch:
diff --git a/larq/callbacks.py b/larq/callbacks.py
--- a/larq/callbacks.py
+++ b/larq/callbacks.py
@@ -4,26 +4,40 @@
 class HyperparameterScheduler(tf.keras.callbacks.Callback):
     """Generic hyperparameter scheduler.
 
+    !!! example
+        ```python
+        bop = lq.optimizers.Bop(threshold=1e-6, gamma=1e-3)
+        adam = tf.keras.optimizers.Adam(0.01)
+        optimizer = lq.optimizers.CaseOptimizer(
+            (lq.optimizers.Bop.is_binary_variable, bop), default_optimizer=adam,
+        )
+        callbacks = [
+            HyperparameterScheduler(lambda x: 0.001 * (0.1 ** (x // 30)), "gamma", bop)
+        ]
+        ```
     # Arguments
+    optimizer: the optimizer that contains the hyperparameter that will be scheduled.
+        Defaults to `self.model.optimizer` if `optimizer == None`.
     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):
+    def __init__(self, schedule, hyperparameter, optimizer=None, verbose=0):
         super(HyperparameterScheduler, self).__init__()
+        self.optimizer = optimizer if optimizer else self.model.optimizer
         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):
+        if not hasattr(self.optimizer, self.hyperparameter):
             raise ValueError(
                 f'Optimizer must have a "{self.hyperparameter}" attribute.'
             )
 
-        hp = getattr(self.model.optimizer, self.hyperparameter)
+        hp = getattr(self.optimizer, self.hyperparameter)
         try:  # new API
             hyperparameter_val = tf.keras.backend.get_value(hp)
             hyperparameter_val = self.schedule(epoch, hyperparameter_val)
@@ -34,10 +48,10 @@ def on_epoch_begin(self, epoch, logs=None):
 
         if self.verbose > 0:
             print(
-                f"Epoch {epoch + 1}: {self.hyperparameter} changning to {tf.keras.backend.get_value(hp)}."
+                f"Epoch {epoch + 1}: {self.hyperparameter} changing 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)
+        hp = getattr(self.optimizer, self.hyperparameter)
         logs[self.hyperparameter] = tf.keras.backend.get_value(hp)


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,4 +1,10 @@
+import numpy as np
 import tensorflow as tf
+from tensorflow.python.keras import testing_utils
+
+import larq as lq
+from larq import testing_utils as lq_testing_utils
+from larq.callbacks import HyperparameterScheduler
 
 
 class LogHistory(tf.keras.callbacks.Callback):
@@ -18,4 +24,76 @@ def on_epoch_end(self, epoch, logs={}):
 
 
 class TestHyperparameterScheduler:
-    pass
+    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 = tf.keras.utils.to_categorical(y_train)
+
+        model = lq_testing_utils.get_small_bnn_model(
+            x_train.shape[1], 20, y_train.shape[1]
+        )
+
+        bop = lq.optimizers.Bop(threshold=1e-6, gamma=1e-3)
+        adam = tf.keras.optimizers.Adam(0.01)
+        case_optimizer = lq.optimizers.CaseOptimizer(
+            (lq.optimizers.Bop.is_binary_variable, bop), default_optimizer=adam,
+        )
+
+        model.compile(
+            loss="categorical_crossentropy",
+            optimizer=case_optimizer,
+            metrics=["accuracy"],
+        )
+
+        def scheduler(x):
+            return 1.0 / (1.0 + x)
+
+        cbk_gamma_scheduler = HyperparameterScheduler(
+            schedule=scheduler,
+            optimizer=model.optimizer.optimizers[0],
+            hyperparameter="gamma",
+            verbose=1,
+        )
+        cbk_threshold_scheduler = HyperparameterScheduler(
+            schedule=scheduler,
+            optimizer=model.optimizer.optimizers[0],
+            hyperparameter="threshold",
+            verbose=1,
+        )
+        cbk_lr_scheduler = HyperparameterScheduler(
+            schedule=scheduler,
+            optimizer=model.optimizer.optimizers[1],
+            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(
+            tf.keras.backend.get_value(model.optimizer.optimizers[0].gamma),
+            scheduler(num_epochs - 1),
+            decimal=8,
+        )
+
+        np.testing.assert_almost_equal(
+            tf.keras.backend.get_value(model.optimizer.optimizers[0].threshold),
+            scheduler(num_epochs - 1),
+            decimal=8,
+        )
+
+        np.testing.assert_almost_equal(
+            tf.keras.backend.get_value(model.optimizer.optimizers[1].lr),
+            scheduler(num_epochs - 1),
+            decimal=8,
+        )