Repository: horovod/horovod

Issue:
tf.keras.optimizers.schedules not supported by distributed optimizer
**Environment:**
1. Framework: (TensorFlow, Keras, PyTorch, MXNet)
tensorflow
2. Framework version:
2.0.0
3. Horovod version:
0.18.2
4. MPI version:
openmpi/3.1.3
5. CUDA version:
cuda/10.0.130
6. NCCL version:
2.4.8
7. Python version:
3.7
8. OS and version:
centos 7
9. GCC version:

**Checklist:**
1. Did you search issues to find if somebody asked this question before?
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)?
4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?

**Bug report:**
Please describe errorneous behavior you're observing and steps to reproduce it.

Using tensorflow's [LearningRateSchedules](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/ExponentialDecay) doesn't appear to work with horovod's distributed optimizer. This might be an easier way of implementing warmup than horovod's keras callbacks.

Modifying the example `tensorflow_keras_mnist.py` with the following lines gives this error message
```python
# Horovod: adjust learning rate based on number of GPUs.
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    0.001 * hvd.size(),
    decay_steps=100000,
    decay_rate=0.96,
    staircase=True)
opt = tf.keras.optimizers.Adam(lr_schedule)

# Horovod: add Horovod DistributedOptimizer.
opt = hvd.DistributedOptimizer(opt)
```

```
[1,0]<stderr>:Traceback (most recent call last):
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 541, in __getattribute__
[1,0]<stderr>:    return super(OptimizerV2, self).__getattribute__(name)
[1,0]<stderr>:AttributeError: 'Adam' object has no attribute 'lr'
[1,0]<stderr>:
[1,0]<stderr>:During handling of the above exception, another exception occurred:
[1,0]<stderr>:
[1,0]<stderr>:Traceback (most recent call last):
[1,0]<stderr>:  File "tensorflow_keras_mnist.py", line 94, in <module>
[1,0]<stderr>:    mnist_model.fit(dataset, steps_per_epoch=500 // hvd.size(), callbacks=callbacks, epochs=24, verbose=verbose)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 728, in fit
[1,0]<stderr>:    use_multiprocessing=use_multiprocessing)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_generator.py", line 693, in fit
[1,0]<stderr>:    steps_name='steps_per_epoch')
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_generator.py", line 196, in model_iteration
[1,0]<stderr>:    callbacks._call_begin_hook(mode)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/callbacks.py", line 249, in _call_begin_hook
[1,0]<stderr>:    self.on_train_begin()
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/callbacks.py", line 365, in on_train_begin
[1,0]<stderr>:    callback.on_train_begin(logs)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/horovod/_keras/callbacks.py", line 137, in on_train_begin
[1,0]<stderr>:    self.initial_lr = self.backend.get_value(self.model.optimizer.lr)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 550, in __getattribute__
[1,0]<stderr>:    return self._get_hyper(name)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 527, in _get_hyper
[1,0]<stderr>:    self._create_hypers()
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 643, in _create_hypers
[1,0]<stderr>:    aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 805, in add_weight
[1,0]<stderr>:    aggregation=aggregation)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 744, in _add_variable_with_custom_getter
[1,0]<stderr>:    **kwargs_for_getter)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer_utils.py", line 139, in make_variable
[1,0]<stderr>:    shape=variable_shape if variable_shape else None)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 258, in __call__
[1,0]<stderr>:    return cls._variable_v1_call(*args, **kwargs)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 219, in _variable_v1_call
[1,0]<stderr>:    shape=shape)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 197, in <lambda>
[1,0]<stderr>:    previous_getter = lambda **kwargs: default_variable_creator(None, **kwargs)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 2507, in default_variable_creator
[1,0]<stderr>:    shape=shape)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 262, in __call__
[1,0]<stderr>:    return super(VariableMetaclass, cls).__call__(*args, **kwargs)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/resource_variable_ops.py", line 1406, in __init__
[1,0]<stderr>:    distribute_strategy=distribute_strategy)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/resourc[1,0]<stderr>:e_variable_ops.py", line 1538, in _init_from_args
[1,0]<stderr>:    name="initial_value", dtype=dtype)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1184, in convert_to_tensor
[1,0]<stderr>:    return convert_to_tensor_v2(value, dtype, preferred_dtype, name)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1242, in convert_to_tensor_v2
[1,0]<stderr>:    as_ref=False)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1296, in internal_convert_to_tensor
[1,0]<stderr>:    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 286, in _constant_tensor_conversion_function
[1,0]<stderr>:    return constant(v, dtype=dtype, name=name)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 227, in constant
[1,0]<stderr>:    allow_broadcast=True)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 235, in _constant_impl
[1,0]<stderr>:    t = convert_to_eager_tensor(value, ctx, dtype)
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/constant_op.py", line 96, in convert_to_eager_tensor
[1,0]<stderr>:    return ops.EagerTensor(value, ctx.device_name, dtype)
[1,0]<stderr>:ValueError: Attempt to convert a value ({'class_name': 'ExponentialDecay', 'config': {'initial_learning_rate': 0.001, 'decay_steps': 100000, 'decay_rate': 0.96, 'staircase': True, 'name': None}}) with an unsupported type (<class 'dict'>) to a Tensor.
[1,0]<stderr>:Exception ignored in: <function _RandomSeedGeneratorDeleter.__del__ at 0x7f7dec944a70>
[1,0]<stderr>:Traceback (most recent call last):
[1,0]<stderr>:  File "/home/pstjohn/miniconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 3009, in __del__
[1,0]<stderr>:AttributeError: 'NoneType' object has no attribute 'device'
--------------------------------------------------------------------------
Primary job  terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------
```



Discussion:

Resolving patch:
diff --git a/horovod/_keras/__init__.py b/horovod/_keras/__init__.py
--- a/horovod/_keras/__init__.py
+++ b/horovod/_keras/__init__.py
@@ -20,17 +20,14 @@
 def create_distributed_optimizer(keras, optimizer, name, device_dense, device_sparse,
                                  compression, sparse_as_dense):
     class _DistributedOptimizer(keras.optimizers.Optimizer):
-        def __init__(self, name, device_dense, device_sparse, compression, sparse_as_dense,
-                     config):
-            if name is None:
-                name = "Distributed%s" % self.__class__.__base__.__name__
-            self._name = name
+        def __init__(self, **kwargs):
+            self._name = name or "Distributed%s" % self.__class__.__base__.__name__
             self._device_dense = device_dense
             self._device_sparse = device_sparse
             self._compression = compression
             self._sparse_as_dense = sparse_as_dense
             self._get_gradients_used = False
-            super(self.__class__, self).__init__(**config)
+            super(self.__class__, self).__init__(**kwargs)
 
         def get_gradients(self, loss, params):
             """
@@ -64,24 +61,19 @@ def get_gradients(self, loss, params):
 
         def apply_gradients(self, *args, **kwargs):
             if not self._get_gradients_used:
-                  raise Exception('`apply_gradients()` was called without a call to '
-                                  '`get_gradients()`. If you\'re using TensorFlow 2.0, '
-                                  'please specify `experimental_run_tf_function=False` in '
-                                  '`compile()`.')
+                raise Exception('`apply_gradients()` was called without a call to '
+                                '`get_gradients()`. If you\'re using TensorFlow 2.0, '
+                                'please specify `experimental_run_tf_function=False` in '
+                                '`compile()`.')
             return super(self.__class__, self).apply_gradients(*args, **kwargs)
 
-        @classmethod
-        def from_config(cls, cfg):
-            return cls(name, device_dense, device_sparse, compression, sparse_as_dense, cfg)
-
     # We dynamically create a new class that inherits from the optimizer that was passed in.
     # The goal is to override get_gradients() method with an allreduce implementation.
     # This class will have the same name as the optimizer it's wrapping, so that the saved
     # model could be easily restored without Horovod.
     cls = type(optimizer.__class__.__name__, (optimizer.__class__,),
                dict(_DistributedOptimizer.__dict__))
-    return cls(name, device_dense, device_sparse, compression, sparse_as_dense,
-               optimizer.get_config())
+    return cls.from_config(optimizer.get_config())
 
 
 def _eval(backend, op_or_result):


Test patch:
diff --git a/test/test_tensorflow2_keras.py b/test/test_tensorflow2_keras.py
new file mode 100644
--- /dev/null
+++ b/test/test_tensorflow2_keras.py
@@ -0,0 +1,100 @@
+# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Tests for horovod.tensorflow.keras."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import numpy as np
+import warnings
+
+from tensorflow import keras
+
+import horovod.tensorflow.keras as hvd
+
+
+class Tf2KerasTests(tf.test.TestCase):
+    """
+    Tests for ops in horovod.tensorflow.keras.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(Tf2KerasTests, self).__init__(*args, **kwargs)
+        warnings.simplefilter('module')
+        hvd.init()
+
+        gpus = tf.config.experimental.list_physical_devices('GPU')
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+        if gpus:
+            tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
+
+    def test_train_model_lr_schedule(self):
+        lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
+            0.001 * hvd.size(),
+            decay_steps=100000,
+            decay_rate=0.96,
+            staircase=True)
+        opt = tf.keras.optimizers.Adam(lr_schedule)
+        opt = hvd.DistributedOptimizer(opt)
+
+        model = keras.models.Sequential()
+        model.add(keras.layers.Dense(2, input_shape=(3,)))
+        model.add(keras.layers.RepeatVector(3))
+        model.add(keras.layers.ThresholdedReLU(0.5))
+        model.compile(loss=keras.losses.mean_squared_error,
+                      optimizer=opt,
+                      metrics=[keras.metrics.categorical_accuracy],
+                      experimental_run_tf_function=False)
+
+        x = np.random.random((1, 3))
+        y = np.random.random((1, 3, 2))
+
+        # No assertions, we just need to verify that it doesn't hang or error
+        callbacks = [hvd.callbacks.BroadcastGlobalVariablesCallback(0)]
+        model.fit(x,
+                  y,
+                  steps_per_epoch=10,
+                  callbacks=callbacks,
+                  epochs=1)
+
+    def test_sparse_as_dense(self):
+        opt = keras.optimizers.RMSprop(lr=0.0001)
+        opt = hvd.DistributedOptimizer(opt, sparse_as_dense=True)
+
+        model = keras.models.Sequential()
+        model.add(keras.layers.Embedding(1000, 64, input_length=10))
+        model.compile(loss=keras.losses.mean_squared_error,
+                      optimizer=opt,
+                      experimental_run_tf_function=False)
+
+        x = np.random.randint(1000, size=(32, 10))
+        y = np.random.random((32, 10, 64))
+        # No assertions, we just need to verify that it doesn't hang
+        model.train_on_batch(x, y)
+
+    def test_from_config(self):
+        opt = keras.optimizers.Adam()
+        hopt = hvd.DistributedOptimizer(opt)
+        cfg = hopt.get_config()
+
+        hopt_copy1 = hopt.from_config(cfg)
+        self.assertEqual(cfg, hopt_copy1.get_config())
+
+        hopt_copy2 = hopt.__class__.from_config(cfg)
+        self.assertEqual(cfg, hopt_copy2.get_config())
diff --git a/test/test_tensorflow_keras.py b/test/test_tensorflow_keras.py
--- a/test/test_tensorflow_keras.py
+++ b/test/test_tensorflow_keras.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Tests for horovod.keras."""
+"""Tests for horovod.tensorflow.keras."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -36,7 +36,7 @@
 
 class TfKerasTests(tf.test.TestCase):
     """
-    Tests for ops in horovod.keras.
+    Tests for ops in horovod.tensorflow.keras.
     """
 
     def __init__(self, *args, **kwargs):