Repository: horovod/horovod

Issue:
hvd.load_model not properly wrapping optimizer in tf.keras 1.15
**Environment:**
1. Framework: tf.keras
2. Framework version: TF 1.15.0
3. Horovod version: 0.19.0
4. MPI version: cray-mpich/7.7.10
5. CUDA version: n/a
6. NCCL version: n/a
7. Python version: 3.7.4
8. OS and version: Cray linux based on SLES 15
9. GCC version: 7.3.0

**Checklist:**
1. Did you search issues to find if somebody asked this question before? Oui
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)? n/a
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)? n/a
4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)? Oui oui

**Bug report:**
Please describe erroneous behavior you're observing and steps to reproduce it.

I've been noticing bad training results when resuming from checkpoint with `hvd.load_model`. I tracked it down to diverging worker models, and noticed that my model optimizers from `hvd.load_model` were not properly wrapped in the horovod DistributedOptimizer. I believe it's a bug in this logic that determines all optimizer classes to wrap right here:
https://github.com/horovod/horovod/blob/d1b13ec131af22b31d0ba999dae15a29991cfeae/horovod/_keras/__init__.py#L113

This part:

    horovod_objects = {
        subclass.__name__.lower(): wrap_optimizer(subclass)
        for subclass in keras.optimizers.Optimizer.__subclasses__()
        if subclass.__module__ == keras.optimizers.Optimizer.__module__
    }

This check on the subclass module matching Optimizer module doesn't work in TF 1.15. E.g., for the SGD optimizer, the class module (the LHS) is actually

    In [9]: tensorflow.python.keras.optimizer_v2.gradient_descent.SGD.__module__
    Out[9]: 'tensorflow.python.keras.optimizer_v2.gradient_descent'

whereas the RHS is

    In [10]: tf.keras.optimizers.Optimizer.__module__
    Out[10]: 'tensorflow.python.keras.optimizer_v2.optimizer_v2’

I have for now implemented a workaround in my code that just removes the module equality comparison and confirm that my optimizers are correctly being wrapped in DistributedOptimizer.

I don't have a suggestion for how to fix this in Horovod. Presumably the module paths are inconsistent across Keras and TF versions :(


Discussion:
Thanks for raising this issue, @sparticlesteve.  I'll take a look this.  Presumably, this means our load_model test is not behaving correctly (or at least missing some edge cases).

Resolving patch:
diff --git a/horovod/_keras/__init__.py b/horovod/_keras/__init__.py
--- a/horovod/_keras/__init__.py
+++ b/horovod/_keras/__init__.py
@@ -109,11 +109,11 @@ def broadcast(backend, value, root_rank, name):
     return _eval(backend, hvd.broadcast(tf.constant(value, name=name), root_rank))
 
 
-def load_model(keras, wrap_optimizer, filepath, custom_optimizers, custom_objects):
+def load_model(keras, wrap_optimizer, optimizer_modules, filepath, custom_optimizers, custom_objects):
     horovod_objects = {
         subclass.__name__.lower(): wrap_optimizer(subclass)
         for subclass in keras.optimizers.Optimizer.__subclasses__()
-        if subclass.__module__ == keras.optimizers.Optimizer.__module__
+        if subclass.__module__ in optimizer_modules
     }
 
     if custom_optimizers is not None:
diff --git a/horovod/keras/__init__.py b/horovod/keras/__init__.py
--- a/horovod/keras/__init__.py
+++ b/horovod/keras/__init__.py
@@ -147,4 +147,5 @@ def load_model(filepath, custom_optimizers=None, custom_objects=None, compressio
     """
     def wrap_optimizer(cls):
         return lambda **kwargs: DistributedOptimizer(cls(**kwargs), compression=compression)
-    return _impl.load_model(keras, wrap_optimizer, filepath, custom_optimizers, custom_objects)
+    optimizer_modules = {keras.optimizers.Optimizer.__module__}
+    return _impl.load_model(keras, wrap_optimizer, optimizer_modules, filepath, custom_optimizers, custom_objects)
diff --git a/horovod/tensorflow/keras/__init__.py b/horovod/tensorflow/keras/__init__.py
--- a/horovod/tensorflow/keras/__init__.py
+++ b/horovod/tensorflow/keras/__init__.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+import inspect
+
 import tensorflow as tf
 
 from distutils.version import LooseVersion
@@ -38,6 +40,15 @@
 from horovod.tensorflow.keras import callbacks
 
 
+try:
+    # In later versions of TensorFlow, optimizers are spread across multiple modules. This set is used to distinguish
+    # stock optimizers that come with tf.keras from custom optimizers that may need to be wrapped specially.
+    _OPTIMIZER_MODULES = set([obj.__module__ for name, obj in inspect.getmembers(tf.keras.optimizers)
+                              if isinstance(obj, type(tf.keras.optimizers.Optimizer))])
+except:
+    _OPTIMIZER_MODULES = set()
+
+
 def DistributedOptimizer(optimizer, name=None,
                          device_dense='', device_sparse='',
                          compression=Compression.none,
@@ -153,5 +164,5 @@ def load_model(filepath, custom_optimizers=None, custom_objects=None, compressio
     """
     def wrap_optimizer(cls):
         return lambda **kwargs: DistributedOptimizer(cls(**kwargs), compression=compression)
-    return _impl.load_model(keras, wrap_optimizer, filepath, custom_optimizers, custom_objects)
+    return _impl.load_model(keras, wrap_optimizer, _OPTIMIZER_MODULES, filepath, custom_optimizers, custom_objects)
 


Test patch:
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
@@ -21,6 +21,7 @@
 
 import tensorflow as tf
 import numpy as np
+import pytest
 import warnings
 
 from distutils.version import LooseVersion
@@ -33,6 +34,8 @@
 
 import horovod.tensorflow.keras as hvd
 
+from common import temppath
+
 
 class TfKerasTests(tf.test.TestCase):
     """
@@ -98,6 +101,185 @@ def test_sparse_as_dense(self):
             # No assertions, we just need to verify that it doesn't hang
             model.train_on_batch(x, y)
 
+    @pytest.mark.skipif(LooseVersion(tf.__version__) < LooseVersion('1.12.0'), reason='TensorFlow version too low')
+    def test_load_model(self):
+        with self.test_session(config=self.config) as sess:
+            K.set_session(sess)
+
+            opt = keras.optimizers.RMSprop(lr=0.0001)
+            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.TimeDistributed(keras.layers.Dense(3)))
+            model.compile(loss=keras.losses.MSE,
+                          optimizer=opt,
+                          metrics=[keras.metrics.categorical_accuracy],
+                          sample_weight_mode='temporal')
+
+            x = np.random.random((1, 3))
+            y = np.random.random((1, 3, 3))
+            model.train_on_batch(x, y)
+
+            with temppath() as fname:
+                model.save(fname)
+
+                new_model = hvd.load_model(fname)
+                new_opt = new_model.optimizer
+
+            self.assertEqual(type(new_opt).__module__, 'horovod._keras')
+            self.assertEqual(type(new_opt).__name__, 'RMSprop')
+            self.assertEqual(K.get_value(opt.lr), K.get_value(new_opt.lr))
+            self._check_optimizer_weights(opt, new_opt)
+
+    @pytest.mark.skipif(LooseVersion(tf.__version__) < LooseVersion('1.12.0'), reason='TensorFlow version too low')
+    def test_load_model_custom_optimizers(self):
+        class TestOptimizer(keras.optimizers.RMSprop):
+            def __init__(self, **kwargs):
+                super(TestOptimizer, self).__init__(**kwargs)
+
+        with self.test_session(config=self.config) as sess:
+            K.set_session(sess)
+
+            opt = TestOptimizer(lr=0.0001)
+            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.TimeDistributed(keras.layers.Dense(3)))
+            model.compile(loss=keras.losses.MSE,
+                          optimizer=opt,
+                          metrics=[keras.metrics.categorical_accuracy],
+                          sample_weight_mode='temporal')
+
+            x = np.random.random((1, 3))
+            y = np.random.random((1, 3, 3))
+            model.train_on_batch(x, y)
+
+            with temppath() as fname:
+                model.save(fname)
+
+                custom_optimizers = [TestOptimizer]
+                new_model = hvd.load_model(fname, custom_optimizers=custom_optimizers)
+                new_opt = new_model.optimizer
+
+            self.assertEqual(type(new_opt).__module__, 'horovod._keras')
+            self.assertEqual(type(new_opt).__name__, 'TestOptimizer')
+            self._check_optimizer_weights(opt, new_opt)
+
+    @pytest.mark.skipif(LooseVersion(tf.__version__) < LooseVersion('1.12.0'), reason='TensorFlow version too low')
+    def test_load_model_custom_objects(self):
+        class TestOptimizer(keras.optimizers.RMSprop):
+            def __init__(self, **kwargs):
+                super(TestOptimizer, self).__init__(**kwargs)
+
+        with self.test_session(config=self.config) as sess:
+            K.set_session(sess)
+
+            opt = TestOptimizer(lr=0.0001)
+            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.TimeDistributed(keras.layers.Dense(3)))
+            model.compile(loss=keras.losses.MSE,
+                          optimizer=opt,
+                          metrics=[keras.metrics.categorical_accuracy],
+                          sample_weight_mode='temporal')
+
+            x = np.random.random((1, 3))
+            y = np.random.random((1, 3, 3))
+            model.train_on_batch(x, y)
+
+            with temppath() as fname:
+                model.save(fname)
+
+                custom_objects = {
+                    'TestOptimizer': lambda **kwargs: hvd.DistributedOptimizer(
+                        TestOptimizer(**kwargs))
+                }
+                new_model = hvd.load_model(fname, custom_objects=custom_objects)
+                new_opt = new_model.optimizer
+
+            self.assertEqual(type(new_opt).__module__, 'horovod._keras')
+            self.assertEqual(type(new_opt).__name__, 'TestOptimizer')
+            self.assertEqual(K.get_value(opt.lr), K.get_value(new_opt.lr))
+            self._check_optimizer_weights(opt, new_opt)
+
+    @pytest.mark.skipif(LooseVersion(tf.__version__) < LooseVersion('1.12.0'), reason='TensorFlow version too low')
+    def test_load_model_broadcast(self):
+        def create_model():
+            opt = keras.optimizers.SGD(lr=0.01 * hvd.size(), momentum=0.9)
+            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.TimeDistributed(keras.layers.Dense(3)))
+            model.compile(loss=keras.losses.MSE,
+                          optimizer=opt,
+                          metrics=[keras.metrics.categorical_accuracy],
+                          sample_weight_mode='temporal')
+
+            return model
+
+        with temppath() as fname:
+            with self.session(config=self.config) as sess:
+                K.set_session(sess)
+
+                model = create_model()
+
+                x = np.random.random((1, 3))
+                y = np.random.random((1, 3, 3))
+                model.train_on_batch(x, y)
+
+                if hvd.rank() == 0:
+                    model.save(fname)
+
+            K.clear_session()
+            with self.session(config=self.config) as sess:
+                K.set_session(sess)
+
+                weight = np.random.random((1, 3))
+
+                if hvd.rank() == 0:
+                    model = hvd.load_model(fname)
+                else:
+                    model = create_model()
+
+                def generator():
+                    while 1:
+                        yield (x, y, weight)
+
+                if hvd.rank() == 0:
+                    self.assertEqual(len(model.optimizer.weights), 5)
+                else:
+                    self.assertEqual(len(model.optimizer.weights), 0)
+
+                # No assertions, we just need to verify that it doesn't hang
+                callbacks = [hvd.callbacks.BroadcastGlobalVariablesCallback(0)]
+                model.fit_generator(generator(),
+                                    steps_per_epoch=1,
+                                    callbacks=callbacks,
+                                    epochs=1,
+                                    verbose=0,
+                                    workers=4,
+                                    initial_epoch=0)
+
+                self.assertEqual(len(model.optimizer.weights), 5)
+
+    def _check_optimizer_weights(self, opt, new_opt):
+        self.assertEqual(len(opt.get_weights()), len(new_opt.get_weights()))
+        for weights, new_weights in zip(opt.get_weights(),
+                                        new_opt.get_weights()):
+            if np.isscalar(weights):
+                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)