(load_path)
| 293 | |
| 294 | |
| 295 | def get_checkpoint_iteration(load_path): |
| 296 | # Read the tracker file and set the iteration. |
| 297 | tracker_filename = get_checkpoint_tracker_filename(load_path) |
| 298 | if not os.path.isfile(tracker_filename): |
| 299 | print_rank_0('WARNING: could not find the metadata file {} '.format( |
| 300 | tracker_filename)) |
| 301 | if os.path.isdir(load_path): |
| 302 | path = os.path.normpath(load_path) |
| 303 | load_dir, tag = os.path.split(path) |
| 304 | print_rank_0('Try to directly load the checkpoint from the directory') |
| 305 | return load_dir, tag, False, True |
| 306 | print_rank_0(' will not load any checkpoints and will start from ' |
| 307 | 'random') |
| 308 | return load_path, 0, False, False |
| 309 | with open(tracker_filename, 'r') as f: |
| 310 | metastring = f.read().strip() |
| 311 | release = metastring == 'release' |
| 312 | # try: |
| 313 | # iteration = int(metastring) |
| 314 | # except ValueError: |
| 315 | # release = metastring == 'release' |
| 316 | # if not release: |
| 317 | # print_rank_0('ERROR: Invalid metadata file {}. Exiting'.format( |
| 318 | # tracker_filename)) |
| 319 | # exit() |
| 320 | |
| 321 | # assert iteration > 0 or release, 'error parsing metadata file {}'.format( |
| 322 | # tracker_filename) |
| 323 | |
| 324 | return load_path, metastring, release, True |
| 325 | |
| 326 | |
| 327 | def load_checkpoint(model, optimizer, lr_scheduler, args, no_deepspeed=False, no_load_optim=False, no_load_rng=False): |
no test coverage detected