Browse Source

!11939 Modify cnnctc network save checkpoint

From: @zhanghuiyao
Reviewed-by: @linqingke,@liangchenghui
Signed-off-by: @linqingke
tags/v1.2.0-rc1
mindspore-ci-bot Gitee 4 years ago
parent
commit
01a0cdf5f0
1 changed files with 6 additions and 3 deletions
  1. +6
    -3
      model_zoo/official/cv/cnnctc/train.py

+ 6
- 3
model_zoo/official/cv/cnnctc/train.py View File

@@ -80,10 +80,13 @@ def train(args_opt, config):
keep_checkpoint_max=config.KEEP_CKPT_MAX_NUM) keep_checkpoint_max=config.KEEP_CKPT_MAX_NUM)
ckpoint_cb = ModelCheckpoint(prefix="CNNCTC", config=config_ck, directory=config.SAVE_PATH) 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: 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__': if __name__ == '__main__':


Loading…
Cancel
Save