From 4a009555a2258cf990d260b768a9db166a82faaa Mon Sep 17 00:00:00 2001 From: chenhaozhe Date: Wed, 21 Apr 2021 14:45:13 +0800 Subject: [PATCH] only save on the first device while training mobilenetv1 --- model_zoo/official/cv/mobilenetv1/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model_zoo/official/cv/mobilenetv1/train.py b/model_zoo/official/cv/mobilenetv1/train.py index aa29eedc46..8da7bed47a 100755 --- a/model_zoo/official/cv/mobilenetv1/train.py +++ b/model_zoo/official/cv/mobilenetv1/train.py @@ -143,7 +143,7 @@ if __name__ == '__main__': time_cb = TimeMonitor(data_size=step_size) loss_cb = LossMonitor() cb = [time_cb, loss_cb] - if config.save_checkpoint: + if config.save_checkpoint and device_id % min(8, args_opt.device_num) == 0: config_ck = CheckpointConfig(save_checkpoint_steps=config.save_checkpoint_epochs * step_size, keep_checkpoint_max=config.keep_checkpoint_max) ckpt_cb = ModelCheckpoint(prefix="mobilenetv1", directory=ckpt_save_dir, config=config_ck)