yoonlee666
5 years ago
1 changed files with 3 additions and 2 deletions
Unified View
Diff Options
+ 3
- 2
model_zoo/official/nlp/bert/run_pretrain.py
View File
| @@ -51,7 +51,7 @@ def run_pretrain(): | |||||
| parser.add_argument("--do_shuffle", type=str, default="true", help="Enable shuffle for dataset, default is true.") | parser.add_argument("--do_shuffle", type=str, default="true", help="Enable shuffle for dataset, default is true.") | ||||
| parser.add_argument("--enable_data_sink", type=str, default="true", help="Enable data sink, default is true.") | parser.add_argument("--enable_data_sink", type=str, default="true", help="Enable data sink, default is true.") | ||||
| parser.add_argument("--data_sink_steps", type=int, default="1", help="Sink steps for each epoch, default is 1.") | parser.add_argument("--data_sink_steps", type=int, default="1", help="Sink steps for each epoch, default is 1.") | ||||
| parser.add_argument("--save_checkpoint_path", type=str, default=None, help="Save checkpoint path") | |||||
| parser.add_argument("--save_checkpoint_path", type=str, default="", help="Save checkpoint path") | |||||
| parser.add_argument("--load_checkpoint_path", type=str, default="", help="Load checkpoint file path") | parser.add_argument("--load_checkpoint_path", type=str, default="", help="Load checkpoint file path") | ||||
| parser.add_argument("--save_checkpoint_steps", type=int, default=1000, help="Save checkpoint steps, " | parser.add_argument("--save_checkpoint_steps", type=int, default=1000, help="Save checkpoint steps, " | ||||
| "default is 1000.") | "default is 1000.") | ||||
| @@ -145,7 +145,8 @@ def run_pretrain(): | |||||
| if args_opt.enable_save_ckpt == "true" and args_opt.device_id % min(8, device_num) == 0: | if args_opt.enable_save_ckpt == "true" and args_opt.device_id % min(8, device_num) == 0: | ||||
| config_ck = CheckpointConfig(save_checkpoint_steps=args_opt.save_checkpoint_steps, | config_ck = CheckpointConfig(save_checkpoint_steps=args_opt.save_checkpoint_steps, | ||||
| keep_checkpoint_max=args_opt.save_checkpoint_num) | keep_checkpoint_max=args_opt.save_checkpoint_num) | ||||
| ckpoint_cb = ModelCheckpoint(prefix='checkpoint_bert', directory=ckpt_save_dir, config=config_ck) | |||||
| ckpoint_cb = ModelCheckpoint(prefix='checkpoint_bert', | |||||
| directory=None if ckpt_save_dir == "" else ckpt_save_dir, config=config_ck) | |||||
| callback.append(ckpoint_cb) | callback.append(ckpoint_cb) | ||||
| if args_opt.load_checkpoint_path: | if args_opt.load_checkpoint_path: | ||||