Fix printing on horovod
@@ -183,6 +183,7 @@ def main():
raise ValueError("\"infer_output_file\" command line parameter is "
"required in inference mode")
infer_config.update(copy.deepcopy(config_module['infer_params']))
+ if hvd is None or hvd.rank() == 0:
deco_print("Inference config:")
pprint.pprint(infer_config)
@@ -201,6 +202,7 @@ def main():
elif 'bench_start' not in train_config:
train_config['bench_start'] = 10 # default value
+ if hvd is None or hvd.rank() == 0:
deco_print("New benchmarking config:")
pprint.pprint(train_config)
args.mode = "train"
@@ -215,6 +217,7 @@ def main():
"Restored checkpoint from {}. Resuming training".format(checkpoint),
)
elif args.mode == 'eval' or args.mode == 'infer':
+ if hvd is None or hvd.rank() == 0:
deco_print("Loading model from {}".format(checkpoint))
with tf.Graph().as_default():
