only save ckpt in rank0 for Transformer

2020-07-06 19:43:49 +08:00 · 2020-07-06 19:43:49 +08:00 · 3ce29513db
parent 3d377c51b9
commit 3ce29513db
1 changed files with 5 additions and 4 deletions
--- a/model_zoo/Transformer/train.py
+++ b/model_zoo/Transformer/train.py
@ -147,6 +147,7 @@ def run_transformer_train():

    callbacks = [TimeMonitor(dataset.get_dataset_size()), LossCallBack()]
    if args.enable_save_ckpt == "true":
+        if device_num == 1 or (device_num > 1 and rank_id == 0):
            ckpt_config = CheckpointConfig(save_checkpoint_steps=args.save_checkpoint_steps,
                                           keep_checkpoint_max=args.save_checkpoint_num)
            ckpoint_cb = ModelCheckpoint(prefix='transformer', directory=args.save_checkpoint_path, config=ckpt_config)