From 177bdd8f4b715a29b9b2fbd861fd092ff1521d75 Mon Sep 17 00:00:00 2001 From: zhanghuiyao <1814619459@qq.com> Date: Mon, 1 Feb 2021 16:35:42 +0800 Subject: [PATCH] Modify cnnctc network save checkpoint. --- model_zoo/official/cv/cnnctc/train.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/model_zoo/official/cv/cnnctc/train.py b/model_zoo/official/cv/cnnctc/train.py index a85f484e41..5f7a487603 100644 --- a/model_zoo/official/cv/cnnctc/train.py +++ b/model_zoo/official/cv/cnnctc/train.py @@ -80,10 +80,13 @@ def train(args_opt, config): keep_checkpoint_max=config.KEEP_CKPT_MAX_NUM) ckpoint_cb = ModelCheckpoint(prefix="CNNCTC", config=config_ck, directory=config.SAVE_PATH) - if args_opt.device_id == 0: - model.train(config.TRAIN_EPOCHS, ds, callbacks=[callback, ckpoint_cb], dataset_sink_mode=False) + if args_opt.run_distribute: + if args_opt.device_id == 0: + model.train(config.TRAIN_EPOCHS, ds, callbacks=[callback, ckpoint_cb], dataset_sink_mode=False) + else: + model.train(config.TRAIN_EPOCHS, ds, callbacks=[callback], dataset_sink_mode=False) else: - model.train(config.TRAIN_EPOCHS, ds, callbacks=[callback], dataset_sink_mode=False) + model.train(config.TRAIN_EPOCHS, ds, callbacks=[callback, ckpoint_cb], dataset_sink_mode=False) if __name__ == '__main__':