Use flags to define collective ops when initializing MirroredStrategy (#6724)

f5073f49 · Haoyu Zhang · GitHub · 0f76239b · f5073f49
Unverified Commit f5073f49 authored May 07, 2019 by Haoyu Zhang Committed by GitHub May 07, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

official/resnet/keras/keras_imagenet_main.py official/resnet/keras/keras_imagenet_main.py +3 -1

No files found.
--- a/official/resnet/keras/keras_imagenet_main.py
+++ b/official/resnet/keras/keras_imagenet_main.py
@@ -125,7 +125,9 @@ def run(flags_obj):
  strategy = distribution_utils.get_distribution_strategy(
      distribution_strategy=flags_obj.distribution_strategy,
      num_gpus=flags_obj.num_gpus,
-      num_workers=distribution_utils.configure_cluster())
+      num_workers=distribution_utils.configure_cluster(),
+      all_reduce_alg=flags_obj.all_reduce_alg,
+      num_packs=flags_obj.num_packs)

  strategy_scope = distribution_utils.get_strategy_scope(strategy)