Fix tests well this time

nvcastet · nvcastet · commit cd5414070b07 · 2023-01-25T09:43:36.000-06:00
Signed-off-by: Nicolas Castet &lt;ncastet@nvidia.com&gt;
diff --git a/examples/tensorflow2/tensorflow2_mnist_data_service_train_fn_compute_side_dispatcher.py b/examples/tensorflow2/tensorflow2_mnist_data_service_train_fn_compute_side_dispatcher.py
@@ -23,9 +23,9 @@
 
 from packaging import version
 if version.parse(tf.keras.__version__.replace("-tf", "+tf")) < version.parse("2.11"):
-    from tensorflow.keras.optimizers import Optimizer
+    from tensorflow.keras import optimizers
 else:
-    from tensorflow.keras.optimizers.legacy import Optimizer
+    from tensorflow.keras.optimizers import legacy as optimizers
 
 # arguments reuse_dataset and round_robin only used when single dispatcher is present
 def train_fn(compute_config: TfDataServiceConfig, reuse_dataset: bool = False, round_robin: bool = False):
@@ -69,7 +69,7 @@ def train_fn(compute_config: TfDataServiceConfig, reuse_dataset: bool = False, r
 
     # Horovod: adjust learning rate based on number of GPUs.
     scaled_lr = 0.001 * hvd.size()
-    opt = Optimizer.Adam(scaled_lr)
+    opt = optimizers.Adam(scaled_lr)
 
     # Horovod: add Horovod DistributedOptimizer.
     opt = hvd.DistributedOptimizer(
diff --git a/test/integration/data/elastic_tensorflow2_main.py b/test/integration/data/elastic_tensorflow2_main.py
@@ -25,9 +25,9 @@
 
 from packaging import version
 if version.parse(tf.keras.__version__.replace("-tf", "+tf")) < version.parse("2.11"):
-    from tensorflow.keras.optimizers import Optimizer
+    from tensorflow.keras import optimizers
 else:
-    from tensorflow.keras.optimizers.legacy import Optimizer
+    from tensorflow.keras.optimizers import legacy as optimizers
 
 parser = argparse.ArgumentParser(description='TensorFlow 2 Elastic Test',
                                  formatter_class=argparse.ArgumentDefaultsHelpFormatter)
@@ -64,7 +64,7 @@
 
 lr = 0.001
 model = tf.keras.Sequential([tf.keras.layers.Dense(2, activation='softmax')])
-optimizer = Optimizer.SGD(lr * hvd.size())
+optimizer = optimizers.SGD(lr * hvd.size())
 
 hostname = os.environ.get('HOROVOD_HOSTNAME')
 start_rank = int(os.environ.get('HOROVOD_RANK', 0))
diff --git a/test/parallel/test_tensorflow2_keras.py b/test/parallel/test_tensorflow2_keras.py
@@ -43,9 +43,9 @@
 
 from packaging import version
 if version.parse(keras.__version__.replace("-tf", "+tf")) < version.parse("2.11"):
-    from keras.optimizers import Optimizer
+    from keras import optimizers
 else:
-    from keras.optimizers.legacy import Optimizer
+    from keras.optimizers import legacy as optimizers
 
 _PRE_TF_2_2_0 = version.parse(tf.__version__) < version.parse("2.2.0")
 
@@ -71,10 +71,7 @@ def __init__(self, *args, **kwargs):
 
     def test_train_model_lr_schedule(self):
         initial_lr = 0.1 * hvd.size()
-        if version.parse(tf.keras.__version__.replace("-tf", "+tf")) < version.parse("2.11"):
-            opt = tf.keras.optimizers.Adam()
-        else:
-            opt = tf.keras.optimizers.legacy.Adam()
+        opt = optimizers.Adam()
         opt = hvd.DistributedOptimizer(opt)
 
         def linear_multiplier(epoch):
@@ -164,10 +161,7 @@ def test_sparse_as_dense(self):
 
     def test_sparse_as_dense_with_grad_aggregation(self):
         backward_passes_per_step = 2
-        if version.parse(keras.__version__.replace("-tf", "+tf")) < version.parse("2.11"):
-            opt = keras.optimizers.RMSprop(lr=0.0001)
-        else:
-            opt = keras.optimizers.legacy.RMSprop(lr=0.0001)
+        opt = optimizers.RMSprop(lr=0.0001)
         opt = hvd.DistributedOptimizer(
             opt,
             sparse_as_dense=True,
@@ -193,10 +187,7 @@ def test_sparse_as_dense_with_grad_aggregation(self):
     def test_grad_aggregation_with_inf_grad(self):
         backward_passes_per_step = 2
         step_count = tf.Variable(0, trainable=False, dtype=tf.int32)
-        if version.parse(tf.keras.__version__.replace("-tf", "+tf")) < version.parse("2.11"):
-            opt = tf.keras.optimizers.SGD()
-        else:
-            opt = tf.keras.optimizers.legacy.SGD()
+        opt = optimizers.SGD()
         opt = hvd.DistributedOptimizer(
             opt,
             backward_passes_per_step=backward_passes_per_step,
@@ -221,10 +212,7 @@ def loss():
         assert tf.math.is_finite(grads_and_vars[0][0])
 
     def test_from_config(self):
-        if version.parse(keras.__version__.replace("-tf", "+tf")) < version.parse("2.11"):
-            opt = keras.optimizers.Adam()
-        else:
-            opt = keras.optimizers.legacy.Adam()
+        opt = optimizers.Adam()
         hopt = hvd.DistributedOptimizer(opt)
         cfg = hopt.get_config()
 
@@ -252,7 +240,7 @@ def test_elastic_state(self):
             [np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
              np.array([0.0, 0.0], dtype=np.float32)])
 
-        optimizer = Optimizer.Adam(0.001 * hvd.size())
+        optimizer = optimizers.Adam(0.001 * hvd.size())
 
         state = hvd.elastic.KerasState(
             model1,
@@ -543,10 +531,7 @@ def test_partial_distributed_optimizer(self):
             model.add(tf.keras.layers.Dense(2, input_shape=(3,), kernel_initializer=initializer, bias_initializer=initializer))
             model.add(tf.keras.layers.RepeatVector(3))
             model.add(tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(3, kernel_initializer=initializer, bias_initializer=initializer)))
-            if version.parse(tf.keras.__version__.replace("-tf", "+tf")) < version.parse("2.11"):
-                opt = tf.keras.optimizers.Adam()
-            else:
-                opt = tf.keras.optimizers.legacy.Adam()
+            opt = optimizers.legacy.Adam()
             model.compile(loss=tf.keras.losses.MSE,
                             metrics=[tf.keras.metrics.categorical_accuracy])
 
