We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `larq/callbacks.py`
Content:
```
1 import tensorflow as tf
2 import numpy as np
3 
4 
5 class QuantizationLogger(tf.keras.callbacks.Callback):
6     """Callback that adds quantization specific metrics.
7 
8     !!! note ""
9         In order for metrics to be picked up by TensorBoard this callback needs to be
10         applied before the TensorBoard callback and use the same update frequency.
11 
12     !!! example
13         ```python
14         callbacks = [QuantizationLogger(), tf.keras.callbacks.TensorBoard()]
15         model.fit(X_train, Y_train, callbacks=callbacks)
16         ```
17 
18     # Metrics
19     - `changed_quantization_ration`: The ration of quantized weights in each layer that
20         changed during the weight update.
21 
22     # Arguments
23     update_freq: `'batch'` or `'epoch'` or integer. When using `'batch'`, computes the
24         metrics after each batch. The same applies for `'epoch'`. If using an integer
25         the callback will compute the metrics every `update_freq` batches.
26         Note that computing too frequently can slow down training.
27     """
28 
29     def __init__(self, update_freq="epoch"):
30         super().__init__()
31         self.batch_previous_weights = {}
32         self.epoch_previous_weights = {}
33         self.update_freq = update_freq if update_freq != "batch" else 1
34         self._quantized_weights = []
35         self._quantized_weight_names = []
36 
37     def set_model(self, model):
38         super().set_model(model)
39         for layer in model.layers:
40             if hasattr(layer, "quantized_weights"):
41                 for i, weight in enumerate(layer.quantized_weights):
42                     self._quantized_weights.append(weight)
43                     self._quantized_weight_names.append(
44                         layer.name if i == 0 else f"{layer.name}_{i}"
45                     )
46 
47     def _maybe_log_and_store(self, storage, logs, should_log=True, should_store=True):
48         if should_log or should_store:
49             values = tf.keras.backend.batch_get_value(self._quantized_weights)
50             for key, value in zip(self._quantized_weight_names, values):
51                 value = value.astype(np.int8)
52                 if should_log:
53                     logs[f"changed_quantization_ration/{key.replace(':', '_')}"] = 1 - (
54                         np.count_nonzero(value == storage[key]) / value.size
55                     )
56                 if should_store:
57                     storage[key] = value
58 
59             if should_log and not should_store:
60                 # We don't need it in the next batch anymore
61                 storage = {}
62 
63     def on_batch_end(self, batch, logs=None):
64         if self.update_freq != "epoch":
65             self._maybe_log_and_store(
66                 self.batch_previous_weights,
67                 logs,
68                 should_log=batch > 0 and (batch + 1) % self.update_freq == 0,
69                 should_store=(batch + 2) % self.update_freq == 0,
70             )
71 
72     def on_train_begin(self, logs=None):
73         self._maybe_log_and_store(self.epoch_previous_weights, logs, should_log=False)
74 
75     def on_epoch_end(self, epoch, logs=None):
76         self._maybe_log_and_store(self.epoch_previous_weights, logs)
77
```

Path: `larq/optimizers_v1.py`
Content:
```
1 import tensorflow as tf
2 import numpy as np
3 import larq as lq
4 
5 from larq import utils
6 from copy import deepcopy
7 
8 
9 @utils.register_keras_custom_object
10 class XavierLearningRateScaling(tf.keras.optimizers.Optimizer):
11     """Optimizer wrapper for Xavier Learning Rate Scaling
12 
13     Scale the weights learning rates respectively with the weights initialization
14 
15     !!! note ""
16         This is a wrapper and does not implement any optimization algorithm.
17 
18     !!! example
19         ```python
20         optimizer = lq.optimizers.XavierLearningRateScaling(
21             tf.keras.optimizers.Adam(0.01), model
22         )
23         ```
24 
25     # Arguments
26     optimizer: A `tf.keras.optimizers.Optimizer`
27     model: A `tf.keras.Model`
28 
29     # References
30     - [BinaryConnect: Training Deep Neural Networks with binary weights during
31       propagations](https://arxiv.org/abs/1511.00363)
32     """
33 
34     def __init__(self, optimizer, model):
35         if int(tf.__version__[0]) == 2:
36             raise NotImplementedError(
37                 "XavierLearningRateScaling is not supported by Tensorflow 2.0."
38             )
39 
40         if not isinstance(optimizer, tf.keras.optimizers.Optimizer):
41             raise ValueError(
42                 f"Expected tf.keras.optimizers.Optimizer, received {type(optimizer)}."
43             )
44         self.optimizer = optimizer
45 
46         if isinstance(model, tf.keras.Model):
47             self.multipliers = {}
48             for layer in model.layers:
49                 if hasattr(layer, "quantized_latent_weights"):
50                     for weight in layer.quantized_latent_weights:
51                         self.multipliers[weight.name] = self.get_lr_multiplier(weight)
52         elif isinstance(model, dict):
53             self.multipliers = model
54         else:
55             raise ValueError(f"Expected tf.keras.Model or dict, received {type(model)}")
56 
57     def get_lr_multiplier(self, weight):
58         shape = weight.get_shape().as_list()
59         n_input = shape[-2]
60         n_output = shape[-1]
61         if len(shape) == 4:
62             kernelsize = np.prod(shape[:-2])
63             coeff = 1.0 / np.sqrt(1.5 / ((kernelsize * (n_input + n_output))))
64         elif len(shape) == 2:
65             coeff = 1.0 / np.sqrt(1.5 / ((1.0 * (n_input + n_output))))
66         else:
67             raise NotImplementedError(
68                 "Xavier Learning rate scaling not implimented for this kernelsize"
69             )
70         return coeff
71 
72     def get_updates(self, loss, params):
73         mult_lr_params = [p for p in params if p.name in self.multipliers]
74         base_lr_params = [p for p in params if p.name not in self.multipliers]
75 
76         updates = []
77         base_lr = self.optimizer.lr
78         for param in mult_lr_params:
79             self.optimizer.lr = base_lr * self.multipliers[param.name]
80             updates.extend(self.optimizer.get_updates(loss, [param]))
81 
82         self.optimizer.lr = base_lr
83         updates.extend(self.optimizer.get_updates(loss, base_lr_params))
84 
85         return updates
86 
87     def __getattr__(self, name):
88         return getattr(self.optimizer, name)
89 
90     def get_config(self):
91         return {
92             "optimizer": {
93                 "class_name": self.optimizer.__class__.__name__,
94                 "config": self.optimizer.get_config(),
95             },
96             "multipliers": self.multipliers,
97         }
98 
99     @classmethod
100     def from_config(cls, config, custom_objects=None):
101         optimizer = tf.keras.optimizers.deserialize(
102             config.pop("optimizer"), custom_objects=custom_objects
103         )
104         return cls(optimizer, config["multipliers"])
105 
106 
107 @utils.register_keras_custom_object
108 class Bop(tf.keras.optimizers.Optimizer):
109     """Binary optimizer (Bop).
110 
111     Bop is a latent-free optimizer for Binarized Neural Networks (BNNs).
112 
113     !!! example
114         ```python
115         optimizer = lq.optimizers.Bop(fp_optimizer=tf.keras.optimizers.Adam(0.01))
116 
117         ```
118 
119     # Arguments
120     fp_optimizer: a `tf.keras.optimizers.Optimizer`.
121     threshold: determines to whether to flip each weight.
122     gamma: the adaptivity rate.
123     name: name of the optimizer.
124 
125     # References
126     - [Latent Weights Do Not Exist: Rethinking Binarized Neural Network Optimization](https://arxiv.org/abs/1906.02107)
127     """
128 
129     def __init__(self, fp_optimizer, threshold=1e-5, gamma=1e-2, name="Bop", **kwargs):
130         super().__init__(**kwargs)
131 
132         if not isinstance(fp_optimizer, tf.keras.optimizers.Optimizer):
133             raise TypeError(
134                 f"Expected tf.keras.optimizers.Optimizer, received {type(fp_optimizer)}."
135             )
136 
137         with tf.keras.backend.name_scope(self.__class__.__name__):
138 
139             self.fp_optimizer = fp_optimizer
140             self.threshold = threshold
141             self.gamma = gamma
142 
143     def get_updates(self, loss, params):
144         grads = self.get_gradients(loss, params)
145         self.updates = []
146 
147         ms = [
148             tf.keras.backend.zeros(
149                 tf.keras.backend.int_shape(p), dtype=tf.keras.backend.dtype(p)
150             )
151             for p in params
152         ]
153 
154         fp_params = []
155 
156         for p, g, m in zip(params, grads, ms):
157             if self.is_binary(p):
158                 m_t = (1 - self.gamma) * m + self.gamma * g
159 
160                 self.updates.append(tf.assign(m, m_t))
161                 self.updates.append(
162                     tf.assign(
163                         p, lq.quantizers.sign(-p * tf.sign(p * m_t - self.threshold))
164                     )
165                 )
166 
167             else:
168                 fp_params.append(p)
169 
170         return self.updates + self.fp_optimizer.get_updates(loss, fp_params)
171 
172     @staticmethod
173     def is_binary(var):
174         return "/kernel" in var.name and "quant_" in var.name
175 
176     def __getattr__(self, name):
177         if name == "lr":
178             return self.fp_optimizer.lr
179         return super().__getattr__(name)
180 
181     def get_config(self):
182         fp_optimizer_config = self.fp_optimizer.get_config()
183         config = {
184             "threshold": self.threshold,
185             "gamma": self.gamma,
186             "fp_optimizer": {
187                 "class_name": self.fp_optimizer.__class__.__name__,
188                 "config": fp_optimizer_config,
189             },
190         }
191         return {**super().get_config(), **config}
192 
193     @classmethod
194     def from_config(cls, config, custom_objects=None):
195         new_config = deepcopy(config)
196         fp_optimizer = tf.keras.optimizers.deserialize(
197             new_config["fp_optimizer"], custom_objects=custom_objects
198         )
199         new_config.pop("fp_optimizer", None)
200         return cls(fp_optimizer, **new_config)
201
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/larq/callbacks.py b/larq/callbacks.py
--- a/larq/callbacks.py
+++ b/larq/callbacks.py
@@ -74,3 +74,45 @@
 
     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 @@
         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 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,

{"golden_diff": "diff --git a/larq/callbacks.py b/larq/callbacks.py\n--- a/larq/callbacks.py\n+++ b/larq/callbacks.py\n@@ -74,3 +74,45 @@\n \n     def on_epoch_end(self, epoch, logs=None):\n         self._maybe_log_and_store(self.epoch_previous_weights, logs)\n+\n+\n+class HyperparameterScheduler(tf.keras.callbacks.Callback):\n+    \"\"\"Generic hyperparameter scheduler.\n+\n+    # Arguments:\n+    schedule: a function that takes an epoch index as input\n+        (integer, indexed from 0) and returns a new hyperparameter as output.\n+    hyperparameter: str. the name of the hyperparameter to be scheduled.\n+    verbose: int. 0: quiet, 1: update messages.\n+    \"\"\"\n+\n+    def __init__(self, schedule, hyperparameter, verbose=0):\n+        super(HyperparameterScheduler, self).__init__()\n+        self.schedule = schedule\n+        self.hyperparameter = hyperparameter\n+        self.verbose = verbose\n+\n+    def on_epoch_begin(self, epoch, logs=None):\n+        if not hasattr(self.model.optimizer, self.hyperparameter):\n+            raise ValueError(\n+                f'Optimizer must have a \"{self.hyperparameter}\" attribute.'\n+            )\n+\n+        hp = getattr(self.model.optimizer, self.hyperparameter)\n+        try:  # new API\n+            hyperparameter_val = tf.keras.backend.get_value(hp)\n+            hyperparameter_val = self.schedule(epoch, hyperparameter_val)\n+        except TypeError:  # Support for old API for backward compatibility\n+            hyperparameter_val = self.schedule(epoch)\n+\n+        tf.keras.backend.set_value(hp, hyperparameter_val)\n+\n+        if self.verbose > 0:\n+            print(\n+                f\"Epoch {epoch + 1}: {self.hyperparameter} changning to {tf.keras.backend.get_value(hp)}.\"\n+            )\n+\n+    def on_epoch_end(self, epoch, logs=None):\n+        logs = logs or {}\n+        hp = getattr(self.model.optimizer, self.hyperparameter)\n+        logs[self.hyperparameter] = tf.keras.backend.get_value(hp)\ndiff --git a/larq/optimizers_v1.py b/larq/optimizers_v1.py\n--- a/larq/optimizers_v1.py\n+++ b/larq/optimizers_v1.py\n@@ -137,8 +137,8 @@\n         with tf.keras.backend.name_scope(self.__class__.__name__):\n \n             self.fp_optimizer = fp_optimizer\n-            self.threshold = threshold\n-            self.gamma = gamma\n+            self.threshold = tf.keras.backend.variable(threshold, name=\"threshold\")\n+            self.gamma = tf.keras.backend.variable(gamma, name=\"gamma\")\n \n     def get_updates(self, loss, params):\n         grads = self.get_gradients(loss, params)\n@@ -181,8 +181,8 @@\n     def get_config(self):\n         fp_optimizer_config = self.fp_optimizer.get_config()\n         config = {\n-            \"threshold\": self.threshold,\n-            \"gamma\": self.gamma,\n+            \"threshold\": float(tf.keras.backend.get_value(self.threshold)),\n+            \"gamma\": float(tf.keras.backend.get_value(self.gamma)),\n             \"fp_optimizer\": {\n                 \"class_name\": self.fp_optimizer.__class__.__name__,\n                 \"config\": fp_optimizer_config,\n", "issue": "Generic hyperparameter scheduler \nThe Idea is to have a TF scheduler which is optimizer/hyperparameter agnostic. \r\nAs an example, it can be used for gamma, threshold or learning rate of the floating point optimizer in Bop.\n", "before_files": [{"content": "import tensorflow as tf\nimport numpy as np\n\n\nclass QuantizationLogger(tf.keras.callbacks.Callback):\n    \"\"\"Callback that adds quantization specific metrics.\n\n    !!! note \"\"\n        In order for metrics to be picked up by TensorBoard this callback needs to be\n        applied before the TensorBoard callback and use the same update frequency.\n\n    !!! example\n        ```python\n        callbacks = [QuantizationLogger(), tf.keras.callbacks.TensorBoard()]\n        model.fit(X_train, Y_train, callbacks=callbacks)\n        ```\n\n    # Metrics\n    - `changed_quantization_ration`: The ration of quantized weights in each layer that\n        changed during the weight update.\n\n    # Arguments\n    update_freq: `'batch'` or `'epoch'` or integer. When using `'batch'`, computes the\n        metrics after each batch. The same applies for `'epoch'`. If using an integer\n        the callback will compute the metrics every `update_freq` batches.\n        Note that computing too frequently can slow down training.\n    \"\"\"\n\n    def __init__(self, update_freq=\"epoch\"):\n        super().__init__()\n        self.batch_previous_weights = {}\n        self.epoch_previous_weights = {}\n        self.update_freq = update_freq if update_freq != \"batch\" else 1\n        self._quantized_weights = []\n        self._quantized_weight_names = []\n\n    def set_model(self, model):\n        super().set_model(model)\n        for layer in model.layers:\n            if hasattr(layer, \"quantized_weights\"):\n                for i, weight in enumerate(layer.quantized_weights):\n                    self._quantized_weights.append(weight)\n                    self._quantized_weight_names.append(\n                        layer.name if i == 0 else f\"{layer.name}_{i}\"\n                    )\n\n    def _maybe_log_and_store(self, storage, logs, should_log=True, should_store=True):\n        if should_log or should_store:\n            values = tf.keras.backend.batch_get_value(self._quantized_weights)\n            for key, value in zip(self._quantized_weight_names, values):\n                value = value.astype(np.int8)\n                if should_log:\n                    logs[f\"changed_quantization_ration/{key.replace(':', '_')}\"] = 1 - (\n                        np.count_nonzero(value == storage[key]) / value.size\n                    )\n                if should_store:\n                    storage[key] = value\n\n            if should_log and not should_store:\n                # We don't need it in the next batch anymore\n                storage = {}\n\n    def on_batch_end(self, batch, logs=None):\n        if self.update_freq != \"epoch\":\n            self._maybe_log_and_store(\n                self.batch_previous_weights,\n                logs,\n                should_log=batch > 0 and (batch + 1) % self.update_freq == 0,\n                should_store=(batch + 2) % self.update_freq == 0,\n            )\n\n    def on_train_begin(self, logs=None):\n        self._maybe_log_and_store(self.epoch_previous_weights, logs, should_log=False)\n\n    def on_epoch_end(self, epoch, logs=None):\n        self._maybe_log_and_store(self.epoch_previous_weights, logs)\n", "path": "larq/callbacks.py"}, {"content": "import tensorflow as tf\nimport numpy as np\nimport larq as lq\n\nfrom larq import utils\nfrom copy import deepcopy\n\n\n@utils.register_keras_custom_object\nclass XavierLearningRateScaling(tf.keras.optimizers.Optimizer):\n    \"\"\"Optimizer wrapper for Xavier Learning Rate Scaling\n\n    Scale the weights learning rates respectively with the weights initialization\n\n    !!! note \"\"\n        This is a wrapper and does not implement any optimization algorithm.\n\n    !!! example\n        ```python\n        optimizer = lq.optimizers.XavierLearningRateScaling(\n            tf.keras.optimizers.Adam(0.01), model\n        )\n        ```\n\n    # Arguments\n    optimizer: A `tf.keras.optimizers.Optimizer`\n    model: A `tf.keras.Model`\n\n    # References\n    - [BinaryConnect: Training Deep Neural Networks with binary weights during\n      propagations](https://arxiv.org/abs/1511.00363)\n    \"\"\"\n\n    def __init__(self, optimizer, model):\n        if int(tf.__version__[0]) == 2:\n            raise NotImplementedError(\n                \"XavierLearningRateScaling is not supported by Tensorflow 2.0.\"\n            )\n\n        if not isinstance(optimizer, tf.keras.optimizers.Optimizer):\n            raise ValueError(\n                f\"Expected tf.keras.optimizers.Optimizer, received {type(optimizer)}.\"\n            )\n        self.optimizer = optimizer\n\n        if isinstance(model, tf.keras.Model):\n            self.multipliers = {}\n            for layer in model.layers:\n                if hasattr(layer, \"quantized_latent_weights\"):\n                    for weight in layer.quantized_latent_weights:\n                        self.multipliers[weight.name] = self.get_lr_multiplier(weight)\n        elif isinstance(model, dict):\n            self.multipliers = model\n        else:\n            raise ValueError(f\"Expected tf.keras.Model or dict, received {type(model)}\")\n\n    def get_lr_multiplier(self, weight):\n        shape = weight.get_shape().as_list()\n        n_input = shape[-2]\n        n_output = shape[-1]\n        if len(shape) == 4:\n            kernelsize = np.prod(shape[:-2])\n            coeff = 1.0 / np.sqrt(1.5 / ((kernelsize * (n_input + n_output))))\n        elif len(shape) == 2:\n            coeff = 1.0 / np.sqrt(1.5 / ((1.0 * (n_input + n_output))))\n        else:\n            raise NotImplementedError(\n                \"Xavier Learning rate scaling not implimented for this kernelsize\"\n            )\n        return coeff\n\n    def get_updates(self, loss, params):\n        mult_lr_params = [p for p in params if p.name in self.multipliers]\n        base_lr_params = [p for p in params if p.name not in self.multipliers]\n\n        updates = []\n        base_lr = self.optimizer.lr\n        for param in mult_lr_params:\n            self.optimizer.lr = base_lr * self.multipliers[param.name]\n            updates.extend(self.optimizer.get_updates(loss, [param]))\n\n        self.optimizer.lr = base_lr\n        updates.extend(self.optimizer.get_updates(loss, base_lr_params))\n\n        return updates\n\n    def __getattr__(self, name):\n        return getattr(self.optimizer, name)\n\n    def get_config(self):\n        return {\n            \"optimizer\": {\n                \"class_name\": self.optimizer.__class__.__name__,\n                \"config\": self.optimizer.get_config(),\n            },\n            \"multipliers\": self.multipliers,\n        }\n\n    @classmethod\n    def from_config(cls, config, custom_objects=None):\n        optimizer = tf.keras.optimizers.deserialize(\n            config.pop(\"optimizer\"), custom_objects=custom_objects\n        )\n        return cls(optimizer, config[\"multipliers\"])\n\n\n@utils.register_keras_custom_object\nclass Bop(tf.keras.optimizers.Optimizer):\n    \"\"\"Binary optimizer (Bop).\n\n    Bop is a latent-free optimizer for Binarized Neural Networks (BNNs).\n\n    !!! example\n        ```python\n        optimizer = lq.optimizers.Bop(fp_optimizer=tf.keras.optimizers.Adam(0.01))\n\n        ```\n\n    # Arguments\n    fp_optimizer: a `tf.keras.optimizers.Optimizer`.\n    threshold: determines to whether to flip each weight.\n    gamma: the adaptivity rate.\n    name: name of the optimizer.\n\n    # References\n    - [Latent Weights Do Not Exist: Rethinking Binarized Neural Network Optimization](https://arxiv.org/abs/1906.02107)\n    \"\"\"\n\n    def __init__(self, fp_optimizer, threshold=1e-5, gamma=1e-2, name=\"Bop\", **kwargs):\n        super().__init__(**kwargs)\n\n        if not isinstance(fp_optimizer, tf.keras.optimizers.Optimizer):\n            raise TypeError(\n                f\"Expected tf.keras.optimizers.Optimizer, received {type(fp_optimizer)}.\"\n            )\n\n        with tf.keras.backend.name_scope(self.__class__.__name__):\n\n            self.fp_optimizer = fp_optimizer\n            self.threshold = threshold\n            self.gamma = gamma\n\n    def get_updates(self, loss, params):\n        grads = self.get_gradients(loss, params)\n        self.updates = []\n\n        ms = [\n            tf.keras.backend.zeros(\n                tf.keras.backend.int_shape(p), dtype=tf.keras.backend.dtype(p)\n            )\n            for p in params\n        ]\n\n        fp_params = []\n\n        for p, g, m in zip(params, grads, ms):\n            if self.is_binary(p):\n                m_t = (1 - self.gamma) * m + self.gamma * g\n\n                self.updates.append(tf.assign(m, m_t))\n                self.updates.append(\n                    tf.assign(\n                        p, lq.quantizers.sign(-p * tf.sign(p * m_t - self.threshold))\n                    )\n                )\n\n            else:\n                fp_params.append(p)\n\n        return self.updates + self.fp_optimizer.get_updates(loss, fp_params)\n\n    @staticmethod\n    def is_binary(var):\n        return \"/kernel\" in var.name and \"quant_\" in var.name\n\n    def __getattr__(self, name):\n        if name == \"lr\":\n            return self.fp_optimizer.lr\n        return super().__getattr__(name)\n\n    def get_config(self):\n        fp_optimizer_config = self.fp_optimizer.get_config()\n        config = {\n            \"threshold\": self.threshold,\n            \"gamma\": self.gamma,\n            \"fp_optimizer\": {\n                \"class_name\": self.fp_optimizer.__class__.__name__,\n                \"config\": fp_optimizer_config,\n            },\n        }\n        return {**super().get_config(), **config}\n\n    @classmethod\n    def from_config(cls, config, custom_objects=None):\n        new_config = deepcopy(config)\n        fp_optimizer = tf.keras.optimizers.deserialize(\n            new_config[\"fp_optimizer\"], custom_objects=custom_objects\n        )\n        new_config.pop(\"fp_optimizer\", None)\n        return cls(fp_optimizer, **new_config)\n", "path": "larq/optimizers_v1.py"}], "after_files": [{"content": "import tensorflow as tf\nimport numpy as np\n\n\nclass QuantizationLogger(tf.keras.callbacks.Callback):\n    \"\"\"Callback that adds quantization specific metrics.\n\n    !!! note \"\"\n        In order for metrics to be picked up by TensorBoard this callback needs to be\n        applied before the TensorBoard callback and use the same update frequency.\n\n    !!! example\n        ```python\n        callbacks = [QuantizationLogger(), tf.keras.callbacks.TensorBoard()]\n        model.fit(X_train, Y_train, callbacks=callbacks)\n        ```\n\n    # Metrics\n    - `changed_quantization_ration`: The ration of quantized weights in each layer that\n        changed during the weight update.\n\n    # Arguments\n    update_freq: `'batch'` or `'epoch'` or integer. When using `'batch'`, computes the\n        metrics after each batch. The same applies for `'epoch'`. If using an integer\n        the callback will compute the metrics every `update_freq` batches.\n        Note that computing too frequently can slow down training.\n    \"\"\"\n\n    def __init__(self, update_freq=\"epoch\"):\n        super().__init__()\n        self.batch_previous_weights = {}\n        self.epoch_previous_weights = {}\n        self.update_freq = update_freq if update_freq != \"batch\" else 1\n        self._quantized_weights = []\n        self._quantized_weight_names = []\n\n    def set_model(self, model):\n        super().set_model(model)\n        for layer in model.layers:\n            if hasattr(layer, \"quantized_weights\"):\n                for i, weight in enumerate(layer.quantized_weights):\n                    self._quantized_weights.append(weight)\n                    self._quantized_weight_names.append(\n                        layer.name if i == 0 else f\"{layer.name}_{i}\"\n                    )\n\n    def _maybe_log_and_store(self, storage, logs, should_log=True, should_store=True):\n        if should_log or should_store:\n            values = tf.keras.backend.batch_get_value(self._quantized_weights)\n            for key, value in zip(self._quantized_weight_names, values):\n                value = value.astype(np.int8)\n                if should_log:\n                    logs[f\"changed_quantization_ration/{key.replace(':', '_')}\"] = 1 - (\n                        np.count_nonzero(value == storage[key]) / value.size\n                    )\n                if should_store:\n                    storage[key] = value\n\n            if should_log and not should_store:\n                # We don't need it in the next batch anymore\n                storage = {}\n\n    def on_batch_end(self, batch, logs=None):\n        if self.update_freq != \"epoch\":\n            self._maybe_log_and_store(\n                self.batch_previous_weights,\n                logs,\n                should_log=batch > 0 and (batch + 1) % self.update_freq == 0,\n                should_store=(batch + 2) % self.update_freq == 0,\n            )\n\n    def on_train_begin(self, logs=None):\n        self._maybe_log_and_store(self.epoch_previous_weights, logs, should_log=False)\n\n    def on_epoch_end(self, epoch, logs=None):\n        self._maybe_log_and_store(self.epoch_previous_weights, logs)\n\n\nclass HyperparameterScheduler(tf.keras.callbacks.Callback):\n    \"\"\"Generic hyperparameter scheduler.\n\n    # Arguments:\n    schedule: a function that takes an epoch index as input\n        (integer, indexed from 0) and returns a new hyperparameter as output.\n    hyperparameter: str. the name of the hyperparameter to be scheduled.\n    verbose: int. 0: quiet, 1: update messages.\n    \"\"\"\n\n    def __init__(self, schedule, hyperparameter, verbose=0):\n        super(HyperparameterScheduler, self).__init__()\n        self.schedule = schedule\n        self.hyperparameter = hyperparameter\n        self.verbose = verbose\n\n    def on_epoch_begin(self, epoch, logs=None):\n        if not hasattr(self.model.optimizer, self.hyperparameter):\n            raise ValueError(\n                f'Optimizer must have a \"{self.hyperparameter}\" attribute.'\n            )\n\n        hp = getattr(self.model.optimizer, self.hyperparameter)\n        try:  # new API\n            hyperparameter_val = tf.keras.backend.get_value(hp)\n            hyperparameter_val = self.schedule(epoch, hyperparameter_val)\n        except TypeError:  # Support for old API for backward compatibility\n            hyperparameter_val = self.schedule(epoch)\n\n        tf.keras.backend.set_value(hp, hyperparameter_val)\n\n        if self.verbose > 0:\n            print(\n                f\"Epoch {epoch + 1}: {self.hyperparameter} changning to {tf.keras.backend.get_value(hp)}.\"\n            )\n\n    def on_epoch_end(self, epoch, logs=None):\n        logs = logs or {}\n        hp = getattr(self.model.optimizer, self.hyperparameter)\n        logs[self.hyperparameter] = tf.keras.backend.get_value(hp)\n", "path": "larq/callbacks.py"}, {"content": "import tensorflow as tf\nimport numpy as np\nimport larq as lq\n\nfrom larq import utils\nfrom copy import deepcopy\n\n\n@utils.register_keras_custom_object\nclass XavierLearningRateScaling(tf.keras.optimizers.Optimizer):\n    \"\"\"Optimizer wrapper for Xavier Learning Rate Scaling\n\n    Scale the weights learning rates respectively with the weights initialization\n\n    !!! note \"\"\n        This is a wrapper and does not implement any optimization algorithm.\n\n    !!! example\n        ```python\n        optimizer = lq.optimizers.XavierLearningRateScaling(\n            tf.keras.optimizers.Adam(0.01), model\n        )\n        ```\n\n    # Arguments\n    optimizer: A `tf.keras.optimizers.Optimizer`\n    model: A `tf.keras.Model`\n\n    # References\n    - [BinaryConnect: Training Deep Neural Networks with binary weights during\n      propagations](https://arxiv.org/abs/1511.00363)\n    \"\"\"\n\n    def __init__(self, optimizer, model):\n        if int(tf.__version__[0]) == 2:\n            raise NotImplementedError(\n                \"XavierLearningRateScaling is not supported by Tensorflow 2.0.\"\n            )\n\n        if not isinstance(optimizer, tf.keras.optimizers.Optimizer):\n            raise ValueError(\n                f\"Expected tf.keras.optimizers.Optimizer, received {type(optimizer)}.\"\n            )\n        self.optimizer = optimizer\n\n        if isinstance(model, tf.keras.Model):\n            self.multipliers = {}\n            for layer in model.layers:\n                if hasattr(layer, \"quantized_latent_weights\"):\n                    for weight in layer.quantized_latent_weights:\n                        self.multipliers[weight.name] = self.get_lr_multiplier(weight)\n        elif isinstance(model, dict):\n            self.multipliers = model\n        else:\n            raise ValueError(f\"Expected tf.keras.Model or dict, received {type(model)}\")\n\n    def get_lr_multiplier(self, weight):\n        shape = weight.get_shape().as_list()\n        n_input = shape[-2]\n        n_output = shape[-1]\n        if len(shape) == 4:\n            kernelsize = np.prod(shape[:-2])\n            coeff = 1.0 / np.sqrt(1.5 / ((kernelsize * (n_input + n_output))))\n        elif len(shape) == 2:\n            coeff = 1.0 / np.sqrt(1.5 / ((1.0 * (n_input + n_output))))\n        else:\n            raise NotImplementedError(\n                \"Xavier Learning rate scaling not implimented for this kernelsize\"\n            )\n        return coeff\n\n    def get_updates(self, loss, params):\n        mult_lr_params = [p for p in params if p.name in self.multipliers]\n        base_lr_params = [p for p in params if p.name not in self.multipliers]\n\n        updates = []\n        base_lr = self.optimizer.lr\n        for param in mult_lr_params:\n            self.optimizer.lr = base_lr * self.multipliers[param.name]\n            updates.extend(self.optimizer.get_updates(loss, [param]))\n\n        self.optimizer.lr = base_lr\n        updates.extend(self.optimizer.get_updates(loss, base_lr_params))\n\n        return updates\n\n    def __getattr__(self, name):\n        return getattr(self.optimizer, name)\n\n    def get_config(self):\n        return {\n            \"optimizer\": {\n                \"class_name\": self.optimizer.__class__.__name__,\n                \"config\": self.optimizer.get_config(),\n            },\n            \"multipliers\": self.multipliers,\n        }\n\n    @classmethod\n    def from_config(cls, config, custom_objects=None):\n        optimizer = tf.keras.optimizers.deserialize(\n            config.pop(\"optimizer\"), custom_objects=custom_objects\n        )\n        return cls(optimizer, config[\"multipliers\"])\n\n\n@utils.register_keras_custom_object\nclass Bop(tf.keras.optimizers.Optimizer):\n    \"\"\"Binary optimizer (Bop).\n\n    Bop is a latent-free optimizer for Binarized Neural Networks (BNNs).\n\n    !!! example\n        ```python\n        optimizer = lq.optimizers.Bop(fp_optimizer=tf.keras.optimizers.Adam(0.01))\n\n        ```\n\n    # Arguments\n    fp_optimizer: a `tf.keras.optimizers.Optimizer`.\n    threshold: determines to whether to flip each weight.\n    gamma: the adaptivity rate.\n    name: name of the optimizer.\n\n    # References\n    - [Latent Weights Do Not Exist: Rethinking Binarized Neural Network Optimization](https://arxiv.org/abs/1906.02107)\n    \"\"\"\n\n    def __init__(self, fp_optimizer, threshold=1e-5, gamma=1e-2, name=\"Bop\", **kwargs):\n        super().__init__(**kwargs)\n\n        if not isinstance(fp_optimizer, tf.keras.optimizers.Optimizer):\n            raise TypeError(\n                f\"Expected tf.keras.optimizers.Optimizer, received {type(fp_optimizer)}.\"\n            )\n\n        with tf.keras.backend.name_scope(self.__class__.__name__):\n\n            self.fp_optimizer = fp_optimizer\n            self.threshold = tf.keras.backend.variable(threshold, name=\"threshold\")\n            self.gamma = tf.keras.backend.variable(gamma, name=\"gamma\")\n\n    def get_updates(self, loss, params):\n        grads = self.get_gradients(loss, params)\n        self.updates = []\n\n        ms = [\n            tf.keras.backend.zeros(\n                tf.keras.backend.int_shape(p), dtype=tf.keras.backend.dtype(p)\n            )\n            for p in params\n        ]\n\n        fp_params = []\n\n        for p, g, m in zip(params, grads, ms):\n            if self.is_binary(p):\n                m_t = (1 - self.gamma) * m + self.gamma * g\n\n                self.updates.append(tf.assign(m, m_t))\n                self.updates.append(\n                    tf.assign(\n                        p, lq.quantizers.sign(-p * tf.sign(p * m_t - self.threshold))\n                    )\n                )\n\n            else:\n                fp_params.append(p)\n\n        return self.updates + self.fp_optimizer.get_updates(loss, fp_params)\n\n    @staticmethod\n    def is_binary(var):\n        return \"/kernel\" in var.name and \"quant_\" in var.name\n\n    def __getattr__(self, name):\n        if name == \"lr\":\n            return self.fp_optimizer.lr\n        return super().__getattr__(name)\n\n    def get_config(self):\n        fp_optimizer_config = self.fp_optimizer.get_config()\n        config = {\n            \"threshold\": float(tf.keras.backend.get_value(self.threshold)),\n            \"gamma\": float(tf.keras.backend.get_value(self.gamma)),\n            \"fp_optimizer\": {\n                \"class_name\": self.fp_optimizer.__class__.__name__,\n                \"config\": fp_optimizer_config,\n            },\n        }\n        return {**super().get_config(), **config}\n\n    @classmethod\n    def from_config(cls, config, custom_objects=None):\n        new_config = deepcopy(config)\n        fp_optimizer = tf.keras.optimizers.deserialize(\n            new_config[\"fp_optimizer\"], custom_objects=custom_objects\n        )\n        new_config.pop(\"fp_optimizer\", None)\n        return cls(fp_optimizer, **new_config)\n", "path": "larq/optimizers_v1.py"}]}