Add the option to run Keras resnet model on multiple workers. (#6368)

official/resnet/keras/keras_imagenet_main.py

@@ -144,7 +144,8 @@ def run(flags_obj):
 
   strategy = distribution_utils.get_distribution_strategy(
       distribution_strategy=flags_obj.distribution_strategy,
-      num_gpus=flags_obj.num_gpus)
+      num_gpus=flags_obj.num_gpus,
+      num_workers=distribution_utils.configure_cluster())
 
   strategy_scope = keras_common.get_strategy_scope(strategy)
 

official/utils/misc/distribution_utils.py

@@ -229,6 +229,8 @@ def configure_cluster(worker_hosts=None, task_index=-1):
   tf_config = json.loads(os.environ.get('TF_CONFIG', '{}'))
   if tf_config:
     num_workers = len(tf_config['cluster']['worker'])
+    if tf_config['cluster'].get('chief', None):
+      num_workers += 1
   elif worker_hosts:
     workers = worker_hosts.split(',')
     num_workers = len(workers)