Repository: larq/larq

Issue:
Bop not working with TF2 on MultiGPU
### Describe the bug
In a multi-gpu environment, using Bop generates the following error:

```
2019-10-11 13:45:47 UTC -- Epoch 1/150
2019-10-11 13:45:50 UTC -- Traceback (most recent call last):
2019-10-11 13:45:50 UTC --   File "/usr/local/bin/nf", line 11, in <module>
2019-10-11 13:45:50 UTC --     load_entry_point('project-final', 'console_scripts', 'nf')()
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/click/core.py", line 764, in __call__
2019-10-11 13:45:50 UTC --     return self.main(*args, **kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/click/core.py", line 717, in main
2019-10-11 13:45:50 UTC --     rv = self.invoke(ctx)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/click/core.py", line 1137, in invoke
2019-10-11 13:45:50 UTC --     return _process_result(sub_ctx.command.invoke(sub_ctx))
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/click/core.py", line 956, in invoke
2019-10-11 13:45:50 UTC --     return ctx.invoke(self.callback, **ctx.params)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/click/core.py", line 555, in invoke
2019-10-11 13:45:50 UTC --     return callback(*args, **kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/zookeeper/cli.py", line 114, in train
2019-10-11 13:45:50 UTC --     function(build_model, dataset, hparams, output_dir, **kwargs)
2019-10-11 13:45:50 UTC --   File "/code/project_final/train.py", line 110, in train
2019-10-11 13:45:50 UTC --     callbacks=callbacks,
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/keras/engine/training.py", line 728, in fit
2019-10-11 13:45:50 UTC --     use_multiprocessing=use_multiprocessing)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/keras/engine/training_v2.py", line 324, in fit
2019-10-11 13:45:50 UTC --     total_epochs=epochs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/keras/engine/training_v2.py", line 123, in run_one_epoch
2019-10-11 13:45:50 UTC --     batch_outs = execution_function(iterator)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 86, in execution_function
2019-10-11 13:45:50 UTC --     distributed_function(input_fn))
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/eager/def_function.py", line 457, in __call__
2019-10-11 13:45:50 UTC --     result = self._call(*args, **kwds)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/eager/def_function.py", line 503, in _call
2019-10-11 13:45:50 UTC --     self._initialize(args, kwds, add_initializers_to=initializer_map)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/eager/def_function.py", line 408, in _initialize
2019-10-11 13:45:50 UTC --     *args, **kwds))
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/eager/function.py", line 1848, in _get_concrete_function_internal_garbage_collected
2019-10-11 13:45:50 UTC --     graph_function, _, _ = self._maybe_define_function(args, kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/eager/function.py", line 2150, in _maybe_define_function
2019-10-11 13:45:50 UTC --     graph_function = self._create_graph_function(args, kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/eager/function.py", line 2041, in _create_graph_function
2019-10-11 13:45:50 UTC --     capture_by_value=self._capture_by_value),
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/framework/func_graph.py", line 915, in func_graph_from_py_func
2019-10-11 13:45:50 UTC --     func_outputs = python_func(*func_args, **func_kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/eager/def_function.py", line 358, in wrapped_fn
2019-10-11 13:45:50 UTC --     return weak_wrapped_fn().__wrapped__(*args, **kwds)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/keras/engine/training_v2_utils.py", line 73, in distributed_function
2019-10-11 13:45:50 UTC --     per_replica_function, args=(model, x, y, sample_weights))
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/distribute_lib.py", line 760, in experimental_run_v2
2019-10-11 13:45:50 UTC --     return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/distribute_lib.py", line 1787, in call_for_each_replica
2019-10-11 13:45:50 UTC --     return self._call_for_each_replica(fn, args, kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/mirrored_strategy.py", line 661, in _call_for_each_replica
2019-10-11 13:45:50 UTC --     fn, args, kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/mirrored_strategy.py", line 196, in _call_for_each_replica
2019-10-11 13:45:50 UTC --     coord.join(threads)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/training/coordinator.py", line 389, in join
2019-10-11 13:45:50 UTC --     six.reraise(*self._exc_info_to_raise)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/six.py", line 693, in reraise
2019-10-11 13:45:50 UTC --     raise value
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/training/coordinator.py", line 297, in stop_on_exception
2019-10-11 13:45:50 UTC --     yield
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/mirrored_strategy.py", line 190, in _call_for_each_replica
2019-10-11 13:45:50 UTC --     **merge_kwargs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 446, in _distributed_apply
2019-10-11 13:45:50 UTC --     ds_reduce_util.ReduceOp.SUM, grads_and_vars)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/distribute_lib.py", line 1481, in batch_reduce_to
2019-10-11 13:45:50 UTC --     return self._batch_reduce_to(reduce_op, value_destination_pairs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/mirrored_strategy.py", line 707, in _batch_reduce_to
2019-10-11 13:45:50 UTC --     value_destination_pairs)
2019-10-11 13:45:50 UTC --   File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/cross_device_ops.py", line 317, in batch_reduce
2019-10-11 13:45:50 UTC --     value_destination_pairs[0][0].values) == 1:
2019-10-11 13:45:50 UTC -- IndexError: list index out of range
```

The code runs fine if I use `tf.keras.optimizers.Adam()` instead of Bop, or if I run it on a single GPU.

### To Reproduce

```
with tf.distribute.MirroredStrategy().scope():
    model = build_model()
    model.compile(
        optimizer=lq.optimizers.Bop(tf.keras.optimizers.Adam()),
        loss="categorical_crossentropy",
        metrics=["categorical_accuracy", "top_k_categorical_accuracy"],
        )
model.fit(train_data)
```

### Environment

TensorFlow version: 2.0.0
Larq version:  0.7.3



Discussion:

Resolving patch:
diff --git a/larq/optimizers_v2.py b/larq/optimizers_v2.py
--- a/larq/optimizers_v2.py
+++ b/larq/optimizers_v2.py
@@ -58,12 +58,16 @@ def _create_slots(self, var_list):
                 self.add_slot(var, "m")
 
     def apply_gradients(self, grads_and_vars, name=None):
-        bin_grads_and_vars = [(g, v) for g, v in grads_and_vars if self.is_binary(v)]
-        fp_grads_and_vars = [(g, v) for g, v in grads_and_vars if not self.is_binary(v)]
+        bin_grads_and_vars, fp_grads_and_vars = [], []
+        for grad, var in grads_and_vars:
+            if self.is_binary(var):
+                bin_grads_and_vars.append((grad, var))
+            else:
+                fp_grads_and_vars.append((grad, var))
 
         bin_train_op = super().apply_gradients(bin_grads_and_vars, name=name)
-        fp_train_op = self.fp_optimizer.apply_gradients(fp_grads_and_vars, name=name)
 
+        fp_train_op = self.fp_optimizer.apply_gradients(fp_grads_and_vars, name=name)
         return tf.group(bin_train_op, fp_train_op, name="train_with_bop")
 
     def _resource_apply_sparse(self, grad, var, indices):


Test patch:
diff --git a/larq/optimizers_test.py b/larq/optimizers_test.py
--- a/larq/optimizers_test.py
+++ b/larq/optimizers_test.py
@@ -23,17 +23,30 @@ def assert_weights(weights, expected):
         np.testing.assert_allclose(np.squeeze(w), e)
 
 
-def _test_optimizer(optimizer, target=0.75, test_kernels_are_binary=True):
+def _test_optimizer(
+    optimizer, target=0.75, test_kernels_are_binary=True, trainable_bn=True
+):
     np.random.seed(1337)
     (x_train, y_train), _ = 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])
+    model = lq_testing_utils.get_small_bnn_model(
+        x_train.shape[1], 20, y_train.shape[1], trainable_bn=trainable_bn
+    )
     model.compile(loss="categorical_crossentropy", optimizer=optimizer, metrics=["acc"])
+
+    initial_vars = [tf.keras.backend.get_value(w) for w in model.trainable_weights]
+
     history = model.fit(x_train, y_train, epochs=2, batch_size=16, verbose=0)
 
+    trained_vars = [tf.keras.backend.get_value(w) for w in model.trainable_weights]
+
+    # check all trainable variables have actually been updated
+    for v0, v1 in zip(initial_vars, trained_vars):
+        assert not np.all(v0 == v1)
+
     # Note that when kernels are treated as latent weights they need not be
     # binary (see https://arxiv.org/abs/1906.02107 for further discussion)
     if test_kernels_are_binary:
@@ -103,6 +116,13 @@ def test_bop_accuracy(self):
             lq.optimizers.Bop(fp_optimizer=tf.keras.optimizers.Adam(0.01)),
             test_kernels_are_binary=True,
         )
+        # test optimizer on model with only binary trainable vars (low accuracy)
+        _test_optimizer(
+            lq.optimizers.Bop(fp_optimizer=tf.keras.optimizers.Adam(0.01)),
+            test_kernels_are_binary=True,
+            trainable_bn=False,
+            target=0,
+        )
 
     @pytest.mark.skipif(
         utils.tf_1_14_or_newer() is False,
diff --git a/larq/testing_utils.py b/larq/testing_utils.py
--- a/larq/testing_utils.py
+++ b/larq/testing_utils.py
@@ -7,7 +7,7 @@
 from tensorflow.python.keras.testing_utils import _thread_local_data, should_run_eagerly
 
 
-def get_small_bnn_model(input_dim, num_hidden, output_dim):
+def get_small_bnn_model(input_dim, num_hidden, output_dim, trainable_bn=True):
     model = tf.keras.models.Sequential()
     model.add(
         lq.layers.QuantDense(
@@ -19,7 +19,7 @@ def get_small_bnn_model(input_dim, num_hidden, output_dim):
             use_bias=False,
         )
     )
-    model.add(tf.keras.layers.BatchNormalization())
+    model.add(tf.keras.layers.BatchNormalization(trainable=trainable_bn))
     model.add(
         lq.layers.QuantDense(
             units=output_dim,