Added back horovod.tensorflow.keras.Compression (horovod#2945)

Signed-off-by: Travis Addair <[email protected]>

Author: tgaddair

CHANGELOG.md

@@ -20,6 +20,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `hvd.tensorflow.keras.Compression`, accidentally removed in v0.22.0. ([#2945](https://github.com/horovod/horovod/pull/2945))
+
 ## [v0.22.0] - 2021-05-18
 
 ### Added

docs/requirements.txt

@@ -1,4 +1,4 @@
-sphinx>=3
+sphinx>=3,<4
 sphinxcontrib-napoleon
 alabaster
 nbsphinx

horovod/keras/__init__.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+import warnings
+
 import keras
 import keras.backend as K
 
@@ -27,7 +29,7 @@
 from horovod.tensorflow import gloo_enabled, gloo_built
 from horovod.tensorflow import nccl_built, ddl_built, ccl_built, cuda_built, rocm_built
 from horovod.tensorflow import Average, Sum
-from horovod.tensorflow import compression
+from horovod.tensorflow.compression import Compression
 
 
 from horovod.keras import callbacks, elastic
@@ -36,7 +38,7 @@
 
 def DistributedOptimizer(optimizer, name=None,
                          device_dense='', device_sparse='',
-                         compression=compression.Compression.none,
+                         compression=Compression.none,
                          sparse_as_dense=False,
                          gradient_predivide_factor=1.0,
                          op=Average,
@@ -165,7 +167,7 @@ def broadcast(value, root_rank, name=None):
     return _impl.broadcast(K, value, root_rank, name)
 
 
-def load_model(filepath, custom_optimizers=None, custom_objects=None, compression=compression.Compression.none):
+def load_model(filepath, custom_optimizers=None, custom_objects=None, compression=Compression.none):
     """
     Loads a saved Keras model with a Horovod DistributedOptimizer.
 

horovod/tensorflow/keras/__init__.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 import inspect
+import warnings
 
 import tensorflow as tf
 
@@ -31,7 +32,7 @@
 from horovod.tensorflow import gloo_enabled, gloo_built
 from horovod.tensorflow import nccl_built, ddl_built, ccl_built, cuda_built, rocm_built
 from horovod.tensorflow import Average, Sum
-from horovod.tensorflow import compression
+from horovod.tensorflow.compression import Compression
 
 
 import horovod._keras as _impl
@@ -49,7 +50,7 @@
 
 def DistributedOptimizer(optimizer, name=None,
                          device_dense='', device_sparse='',
-                         compression=compression.Compression.none,
+                         compression=Compression.none,
                          sparse_as_dense=False,
                          gradient_predivide_factor=1.0,
                          op=Average,
@@ -149,7 +150,7 @@ def allreduce(value, name=None, average=None,
               prescale_factor=1.0,
               postscale_factor=1.0,
               op=None,
-              compression=compression.Compression.none):
+              compression=Compression.none):
     """
     Perform an allreduce on a tensor-compatible value.
 
@@ -210,7 +211,7 @@ def broadcast(value, root_rank, name=None):
     return _impl.broadcast(K, value, root_rank, name)
 
 
-def load_model(filepath, custom_optimizers=None, custom_objects=None, compression=compression.Compression.none):
+def load_model(filepath, custom_optimizers=None, custom_objects=None, compression=Compression.none):
     """
     Loads a saved Keras model with a Horovod DistributedOptimizer.