| 54 | |
| 55 | |
| 56 | def save_checkpoint(items, step, model_dir, ignore=[], |
| 57 | keep_every_n=10000000): |
| 58 | if not os.path.exists(model_dir): |
| 59 | os.makedirs(model_dir) |
| 60 | path_without_step = os.path.join(model_dir, 'model_checkpoint') |
| 61 | step_padded = format(step, '08d') |
| 62 | state_dict = items["model"].state_dict() |
| 63 | if ignore: |
| 64 | for key in state_dict.keys(): |
| 65 | for item in ignore: |
| 66 | if key.startswith(item): |
| 67 | state_dict.pop(key) |
| 68 | path_with_step = '{}-{}'.format(path_without_step, step_padded) |
| 69 | |
| 70 | saved_dic = {} |
| 71 | for key in items: |
| 72 | saved_dic[key] = items[key].state_dict() |
| 73 | torch.save({**saved_dic, "step": step}, path_with_step) |
| 74 | |
| 75 | try: |
| 76 | os.unlink(path_without_step) |
| 77 | except FileNotFoundError: |
| 78 | pass |
| 79 | try: |
| 80 | os.symlink(os.path.basename(path_with_step), path_without_step) |
| 81 | except OSError: |
| 82 | shutil.copy2(path_with_step, path_without_step) |
| 83 | |
| 84 | # Cull old checkpoints. |
| 85 | if keep_every_n is not None: |
| 86 | all_checkpoints = [] |
| 87 | for name in os.listdir(model_dir): |
| 88 | m = CHECKPOINT_PATTERN.match(name) |
| 89 | if m is None or name == os.path.basename(path_with_step): |
| 90 | continue |
| 91 | checkpoint_step = int(m.group(1)) |
| 92 | all_checkpoints.append((checkpoint_step, name)) |
| 93 | all_checkpoints.sort() |
| 94 | |
| 95 | last_step = float('-inf') |
| 96 | for checkpoint_step, name in all_checkpoints: |
| 97 | if checkpoint_step - last_step >= keep_every_n: |
| 98 | last_step = checkpoint_step |
| 99 | continue |
| 100 | os.unlink(os.path.join(model_dir, name)) |
| 101 | |
| 102 | |
| 103 | class Saver(object): |