fix estimator train_and_evaluate in multi-worker multi-gpu setting (#23)

Author: SeaOfOcean

epl/__init__.py

@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow.python.eager.context import executing_eagerly
+
 from epl.cluster import Cluster
 from epl.env import Env
 from epl.config import Config
@@ -35,6 +37,9 @@
 
 def init(config=None):
   """Init EPL."""
+  if executing_eagerly():
+    raise RuntimeError("Tensorflow eager mode is not supported by EPL now, " + \
+                       "please do not call tf.enable_eager_execution()")
   env = Env.get()
   env.reset()
   env.init(config)

epl/parallel/hooks.py

@@ -39,6 +39,7 @@
 from tensorflow.python.estimator import estimator
 from tensorflow.python.estimator import run_config
 from tensorflow.python.estimator import training
+from tensorflow.python.estimator.model_fn import ModeKeys
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import device as pydev
 from tensorflow.python.framework import dtypes
@@ -913,17 +914,23 @@ def hook_call_model_fn(self, *args, **kwargs):
 
 def _sync_signal():
   """Sync a tensor among constructor workers."""
-  all_devices = Env.get().cluster.virtual_devices[0].all_devices
-  all_devices = [device for device in all_devices if 'GPU:0' in device]
+  vd = Env.get().cluster.virtual_devices[0]
+  all_devices = vd.all_devices
+  local_devices = vd.local_devices
   if len(all_devices) <= 1:
     return
+
   with ops.Graph().as_default():
-    signal = constant_op.constant([1])
-    with ops.device(Env.get().cluster.current_worker_chief_gpu()):
-      comm = create_serial_communicator(name="BROADCAST_SIGNAL", devices=all_devices)
-      sync_tensor = comm.broadcast([signal])
-    with monitored_session.MonitoredTrainingSession() as sess:
-      sess.run(sync_tensor)
+    sync_ops = []
+    comm = create_serial_communicator(name="BROADCAST_SIGNAL", devices=all_devices)
+    for local_device in local_devices:
+      with ops.device(local_device):
+        signal = constant_op.constant([1])
+        Graph.get().set_model_mode(ModeKeys.EVAL)
+        sync_tensor = comm.broadcast([signal])
+        sync_ops.append(sync_tensor)
+    with session.Session() as sess:
+      sess.run(sync_ops)
 
 
 def estimator_training_evaluate_hook():

epl/utils/launcher.py

@@ -152,7 +152,8 @@ def get_run_script(args):
     script += "export GPU_STATUS_FILE={}/GPU_STATUS_{}.json\n".format(task_dir, index)
     script += get_tf_config(args.num_workers, index, ports, args.machine_list)
     script += ' CUDA_VISIBLE_DEVICES=' + get_visible_devices(args.gpu_per_worker, offset)
-    script += ' bash {}'.format(args.training_script) + ' ' + task_dir
+    cmd = 'python' if args.training_script.endswith('.py') else 'bash'
+    script += ' {} {}'.format(cmd, args.training_script) + ' ' + task_dir
     if args.debug and index == 0:
       print("Enable debug mode")
     else:

tests/Makefile

@@ -11,6 +11,7 @@ test:
 	CUDA_VISIBLE_DEVICES='' ./launch.sh header_test.py
 	PYTHONPATH=../ $(PYTHON) -m epl.utils.launcher --num_workers=2 --gpu_per_worker=1 --debug=True test_launcher.sh
 	PYTHONPATH=../ $(PYTHON) -m epl.utils.launcher --num_workers=2 --gpu_per_worker=1 --debug=True test_amp_parallel.sh
+	CUDA_VISIBLE_DEVICES=$(GPU_1) ./launch.sh eager_test.py
 	CUDA_VISIBLE_DEVICES=$(GPU_4) ./launch.sh strategy_new_test.py
 	CUDA_VISIBLE_DEVICES=$(GPU_4) ./launch.sh auto_cluster_test.py
 	CUDA_VISIBLE_DEVICES=$(GPU_2) ./launch.sh estimator_test.py

tests/eager_test.py

@@ -0,0 +1,43 @@
+# Copyright 2021 Alibaba Group Holding Limited. 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.
+# =============================================================================
+"""Eager mode test."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import tensorflow as tf
+from tensorflow.python.platform import test
+
+import epl
+
+
+# pylint: disable=missing-docstring,unused-variable
+# pylint: disable=protected-access
+class EagerTest(test.TestCase):
+  def test_eager_error(self):
+    assert tf.executing_eagerly() is True
+    with self.assertRaises(RuntimeError) as ctx:
+      epl.init()
+    self.assertTrue("Tensorflow eager mode is not supported by EPL now" in str(ctx.exception))
+
+
+# pylint: enable=missing-docstring,unused-variable
+# pylint: enable=protected-access
+
+if __name__ == "__main__":
+  tf.enable_eager_execution()
+  test.main()

tests/estimator_dp_example.py

@@ -0,0 +1,98 @@
+# Copyright 2021 Alibaba Group Holding Limited. 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.
+# =============================================================================
+"""Test for estimator API."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import shutil
+import tempfile
+import numpy as np
+
+import tensorflow as tf
+
+import epl
+from test_utils import fix_randomness
+
+
+# pylint: disable=missing-docstring,unused-variable
+# pylint: disable=protected-access
+fix_randomness()
+def model_fn(features, labels, mode):
+  """Model function."""
+  with tf.variable_scope('lr_softmax'):
+    weights = tf.get_variable('weights', initializer=tf.zeros([78, 10]))
+    biases = tf.get_variable('biases', initializer=tf.zeros([10]))
+    logits = tf.matmul(features, weights) + biases
+  global_step = tf.train.get_or_create_global_step()
+  loss = tf.losses.sparse_softmax_cross_entropy(labels=labels,
+                                                logits=logits,
+                                                weights=1.0)
+  if mode == tf.estimator.ModeKeys.TRAIN:
+    opt = tf.train.AdamOptimizer(0.1, name='adam')
+    train_op = opt.minimize(loss, global_step=global_step, name='train')
+    if epl.env.Env.get().cluster:
+      epl.add_to_collection(loss, epl.GraphKeys.GLOBAL_MEAN_OBJECTS)
+    return tf.estimator.EstimatorSpec(
+        mode=mode,
+        loss=loss,
+        train_op=train_op)
+  elif mode == tf.estimator.ModeKeys.EVAL:
+    return tf.estimator.EstimatorSpec(
+        mode=mode,
+        loss=loss)
+  else:
+    raise ValueError(
+        "Only TRAIN and EVAL modes are supported: %s" % (mode))
+
+
+def model_fn_replicate(features, labels, mode):
+  with epl.replicate(device_count=1):
+    return model_fn(features, labels, mode)
+
+def input_fn(batch_size=2, estimator=True):
+  """input function."""
+  num_x = np.random.randint(0, 10, (200, 78)).astype(dtype=np.float32)
+  num_y = np.random.randint(0, 10, 200).astype(dtype=np.int32)
+  dataset = tf.data.Dataset.from_tensor_slices((num_x, num_y)) \
+      .batch(batch_size).repeat(10)
+  if not estimator:
+    iterator = dataset.make_initializable_iterator()
+    tf.add_to_collection(tf.GraphKeys.TABLE_INITIALIZERS, iterator.initializer)
+  return dataset
+
+def _main():
+  config = epl.Config()
+  epl.init(config)
+  model_dir = tempfile.mkdtemp()
+  model_dir = os.path.join(model_dir, str(epl.Env.get().cluster.worker_index))
+  run_config = tf.estimator.RunConfig(save_checkpoints_steps=5,
+                                      model_dir=model_dir)
+  estimator = tf.estimator.Estimator(model_fn=model_fn_replicate, config=run_config)
+  hooks = []
+  hooks.append(tf.train.StepCounterHook(every_n_steps=1))
+  train_spec = tf.estimator.TrainSpec(input_fn=input_fn, max_steps=30,
+                                      hooks=hooks)
+  eval_spec = tf.estimator.EvalSpec(input_fn=input_fn, steps=10)
+  tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
+  shutil.rmtree(model_dir)
+# pylint: enable=missing-docstring,unused-variable
+# pylint: enable=protected-access
+
+if __name__ == "__main__":
+  tf.logging.set_verbosity(tf.logging.INFO)
+  _main()