r""" Load checkpoint process.
(args_param, sink_size, dataset, model, network, epoch)
| 331 | |
| 332 | |
| 333 | def restore_checkpoint(args_param, sink_size, dataset, model, network, epoch): |
| 334 | r""" |
| 335 | Load checkpoint process. |
| 336 | """ |
| 337 | print("======start single checkpoint", flush=True) |
| 338 | ckpt_name = args_param.ckpt_name_prefix |
| 339 | ckpt_pattern = os.path.join(args_param.save_checkpoint_path, "rank_{}".format(D.get_rank()), f"{ckpt_name}*.ckpt") |
| 340 | ckpt_all_files = glob.glob(ckpt_pattern) |
| 341 | |
| 342 | if not ckpt_all_files: |
| 343 | print(f"There is no ckpt file in {args_param.save_checkpoint_path}, " |
| 344 | f"current ckpt_files found is {ckpt_all_files} " |
| 345 | f"with pattern {ckpt_pattern}, so skip the loading.") |
| 346 | return |
| 347 | |
| 348 | ckpt_exp_pattern = os.path.join(args_param.save_checkpoint_path, "rank_{}".format(D.get_rank()), |
| 349 | f"{ckpt_name}*_breakpoint.ckpt") |
| 350 | ckpt_exp_files = glob.glob(ckpt_exp_pattern) |
| 351 | ckpt_files = [] |
| 352 | for file in ckpt_all_files: |
| 353 | if file not in ckpt_exp_files: |
| 354 | ckpt_files.append(file) |
| 355 | |
| 356 | if not ckpt_files: |
| 357 | print(f"There is no ckpt file in {args_param.save_checkpoint_path}, " |
| 358 | f"current ckpt_files found is {ckpt_files} " |
| 359 | f"with pattern {ckpt_pattern}, so skip the loading.") |
| 360 | return |
| 361 | |
| 362 | ckpt_files.sort(key=os.path.getmtime, reverse=True) |
| 363 | time_stamp = datetime.datetime.now() |
| 364 | print(f"time stamp {time_stamp.strftime('%Y.%m.%d-%H:%M:%S')} pre trained ckpt model {ckpt_files} loading", |
| 365 | flush=True) |
| 366 | # Load checkpoint files latest file |
| 367 | print(f'Start to load from {ckpt_files[0]}') |
| 368 | param_dict = load_checkpoint(ckpt_files[0]) |
| 369 | if param_dict.get("epoch_num") and param_dict.get("step_num"): |
| 370 | args_param.has_trained_epoches = int(param_dict["epoch_num"].data.asnumpy()) |
| 371 | args_param.has_trained_steps = int(param_dict["step_num"].data.asnumpy()) |
| 372 | model.build(train_dataset=dataset, sink_size=sink_size, epoch=epoch) |
| 373 | load_param_into_net(network, param_dict) |
| 374 | |
| 375 | |
| 376 | def get_exception_checkpoints(args_param): |
no test coverage detected