Repository: horovod/horovod

Issue:
DistributedOptimzer is not compatible with keras.Optimizer
**Environment:**
1. Framework: (TensorFlow, Keras, PyTorch, MXNet)
2. Framework version: 1.14
3. Horovod version: 0.16.4
7. Python version: 3.6.8

**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.
Horovod DistributedOptimzer wrapper is not compatible with keras:



```
import tensorflow as tf
import horovod.tensorflow.keras as hvd

hvd.init()
opt = tf.keras.optimizers.Adam()
hopt = hvd.DistributedOptimizer(opt)
opt.get_config()
cfg = hopt.get_config()
opt_copy = opt.from_config(cfg)
opt_copy = opt.__class__.from_config(cfg)
hopt_copy = hopt.from_config(cfg) # TypeError: __init__() got an unexpected keyword argument 'learning_rate'
hopt_copy = hopt.__class__.from_config(cfg) # TypeError: __init__() got an unexpected keyword argument 'learning_rate'
```



Discussion:
Hey @yselivonchyk, for using `tf.keras` can you try importing `horovod.tensorflow.keras` instead of `horovod.tensorflow`?
@tgaddair Thank you for pointing that out. I updated the issue to reflect how hvd.tensorflow.keras is not compatible with tf.keras.optimizers.Optimizer.
To put it in context, this functionality is used by `tensorflow.python.keras.models.clone_and_build_model` which is called by` tf.keras.experimental.export_saved_model`
Hey @yselivonchyk, for the new example, can you try this instead:

```
import tensorflow as tf
import horovod.tensorflow.keras as hvd

hvd.init()
opt = tf.keras.optimizers.Adam()
hopt = hvd.DistributedOptimizer(opt)
opt.get_config()
cfg = hopt.get_config()
opt_copy = opt.from_config(cfg)
opt_copy = opt.__class__.from_config(cfg)
hopt_copy = hvd.DistributedOptimizer(opt_copy)
```

Basically, we never directly instantiate the `DistributedOptimizer`, we always create it by wrapping another optimizer.

Can you give me an example of how you're using `tf.keras.experimental.export_saved_model`?  Maybe we can workaround this by exporting the non-distributed optimizer instead.
@tgaddair 

I am using it with

```
model.fit(...)
if hvd.rank()==0:
    tf.contrib.saved_model.save_keras_model(model, args.model_dir)
```

Where ` tf.contrib.saved_model.save_keras_model `maps into `tf.keras.experimental.export_saved_model` and than `orig_optimizer.__class__.from_config(optimizer_config)` inside `clone_and_build_model()` fails with the error `# TypeError: __init__() got an unexpected keyword argument 'learning_rate'`.

To get around it I had to rewrite `clone_and_build_model` adding next code snippet:

```
      if "horovod._keras" not in str(type(orig_optimizer)):
        optimizer = orig_optimizer.__class__.from_config(optimizer_config)
      else:
        optimizer = orig_optimizer.__class__.__bases__[0].from_config(optimizer_config)
```

I would assume, that for horovod optimzer a reasonable code update would be to add next method to DistributedOptimizer (https://github.com/horovod/horovod/blob/master/horovod/_keras/__init__.py)

```
def from_config(cfg):
   return _DistributedOptimizer(optimizer.from_config(cfg))
```

Resolving patch:
diff --git a/horovod/_keras/__init__.py b/horovod/_keras/__init__.py
--- a/horovod/_keras/__init__.py
+++ b/horovod/_keras/__init__.py
@@ -70,6 +70,10 @@ def apply_gradients(self, *args, **kwargs):
                                   '`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
diff --git a/horovod/common/util.py b/horovod/common/util.py
--- a/horovod/common/util.py
+++ b/horovod/common/util.py
@@ -15,7 +15,6 @@
 # =============================================================================
 
 from contextlib import contextmanager
-import importlib
 from multiprocessing import Process, Queue
 import os
 import sysconfig


Test patch:
diff --git a/test/test_keras.py b/test/test_keras.py
--- a/test/test_keras.py
+++ b/test/test_keras.py
@@ -238,3 +238,17 @@ def _check_optimizer_weights(self, opt, new_opt):
                 self.assertEqual(weights, new_weights)
             else:
                 self.assertListEqual(weights.tolist(), new_weights.tolist())
+
+    def test_from_config(self):
+        with self.test_session(config=self.config) as sess:
+            K.set_session(sess)
+
+            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
@@ -97,3 +97,17 @@ def test_sparse_as_dense(self):
             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):
+        with self.test_session(config=self.config) as sess:
+            K.set_session(sess)
+
+            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())