Resume if checkpoint is available. Return - epoch of loaded checkpoint.
(checkpoint_dir, model_no_ddp, optimizer)
| 32 | |
| 33 | |
| 34 | def resume_if_possible(checkpoint_dir, model_no_ddp, optimizer): |
| 35 | """ |
| 36 | Resume if checkpoint is available. |
| 37 | Return |
| 38 | - epoch of loaded checkpoint. |
| 39 | """ |
| 40 | epoch = -1 |
| 41 | best_val_metrics = {} |
| 42 | if not os.path.isdir(checkpoint_dir): |
| 43 | return epoch, best_val_metrics |
| 44 | |
| 45 | last_checkpoint = os.path.join(checkpoint_dir, "checkpoint.pth") |
| 46 | if not os.path.isfile(last_checkpoint): |
| 47 | return epoch, best_val_metrics |
| 48 | |
| 49 | sd = torch.load(last_checkpoint, map_location=torch.device("cpu")) |
| 50 | epoch = sd["epoch"] |
| 51 | best_val_metrics = sd["best_val_metrics"] |
| 52 | print(f"Found checkpoint at {epoch}. Resuming.") |
| 53 | |
| 54 | model_no_ddp.load_state_dict(sd["model"], strict=False) |
| 55 | try: |
| 56 | optimizer.load_state_dict(sd["optimizer"]) |
| 57 | except: |
| 58 | print('optimizer weights could not be loaded') |
| 59 | print( |
| 60 | f"Loaded model and optimizer state at {epoch}. Loaded best val metrics so far." |
| 61 | ) |
| 62 | return epoch, best_val_metrics |