Unverified Commit 680eb35c authored by Yuefeng Zhou's avatar Yuefeng Zhou Committed by GitHub
Browse files

Allow distibution_utils.py to worker with PSStrategy or none strategy (#7135)

when there are multiple workers.
parent 58a3de6c
...@@ -114,13 +114,13 @@ def get_distribution_strategy(distribution_strategy="default", ...@@ -114,13 +114,13 @@ def get_distribution_strategy(distribution_strategy="default",
distribution_strategy = distribution_strategy.lower() distribution_strategy = distribution_strategy.lower()
if distribution_strategy == "off": if distribution_strategy == "off":
if num_gpus > 1 or num_workers > 1: if num_gpus > 1:
raise ValueError( raise ValueError(
"When {} GPUs and {} workers are specified, distribution_strategy " "When {} GPUs and {} workers are specified, distribution_strategy "
"flag cannot be set to 'off'.".format(num_gpus, num_workers)) "flag cannot be set to 'off'.".format(num_gpus, num_workers))
return None return None
if distribution_strategy == "multi_worker_mirrored" or num_workers > 1: if distribution_strategy == "multi_worker_mirrored":
return tf.distribute.experimental.MultiWorkerMirroredStrategy( return tf.distribute.experimental.MultiWorkerMirroredStrategy(
communication=_collective_communication(all_reduce_alg)) communication=_collective_communication(all_reduce_alg))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment