Change `CollectiveAllReduceStrategy` to `MultiWorkerMirroredStrategy`. (#6282)

* s/CollectiveAllReduceStrategy/MultiWorkerMirroredStrategy * More s/contrib.distribute/distribute.experimental

Change `CollectiveAllReduceStrategy` to `MultiWorkerMirroredStrategy`. (#6282)
* s/CollectiveAllReduceStrategy/MultiWorkerMirroredStrategy * More s/contrib.distribute/distribute.experimental
d793ea82 · Ayush Dubey · Toby Boyd · 54dffe2e · d793ea82
Commit d793ea82 authored Feb 28, 2019 by Ayush Dubey Committed by Toby Boyd Feb 28, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 4 deletions

official/utils/misc/distribution_utils.py official/utils/misc/distribution_utils.py +2 -4

No files found.
--- a/official/utils/misc/distribution_utils.py
+++ b/official/utils/misc/distribution_utils.py
@@ -63,8 +63,7 @@ def get_distribution_strategy(distribution_strategy="default",
    return None
  if distribution_strategy == "multi_worker_mirrored" or num_workers > 1:
-    return tf.contrib.distribute.CollectiveAllReduceStrategy(
+    return tf.distribute.experimental.MultiWorkerMirroredStrategy()
-        num_gpus_per_worker=num_gpus)
  if (distribution_strategy == "one_device" or
      (distribution_strategy == "default" and num_gpus <= 1)):
@@ -91,8 +90,7 @@ def get_distribution_strategy(distribution_strategy="default",
      return tf.distribute.MirroredStrategy(devices=devices)
  if distribution_strategy == "parameter_server":
-    return tf.contrib.distribute.ParameterServerStrategy(
+    return tf.distribute.experimental.ParameterServerStrategy()
-        num_gpus_per_worker=num_gpus)
  raise ValueError(
      "Unrecognized Distribution Strategy: %r" % distribution_strategy)