MCPcopy Create free account
hub / github.com/zai-org/CodeGeeX / restore_checkpoint

Function restore_checkpoint

codegeex/mindspore/train.py:333–373  ·  view source on GitHub ↗

r""" Load checkpoint process.

(args_param, sink_size, dataset, model, network, epoch)

Source from the content-addressed store, hash-verified

331
332
333def 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
376def get_exception_checkpoints(args_param):

Callers 1

run_train_pipelineFunction · 0.85

Calls 2

load_checkpointFunction · 0.85
getMethod · 0.45

Tested by

no test coverage detected