| 40 | dist.destroy_process_group() |
| 41 | |
| 42 | def save_checkpoint(args, model, optimizer=None, step=None, epoch=None): |
| 43 | CHECKPOINT_PATH = None |
| 44 | if step is not None: |
| 45 | CHECKPOINT_PATH = os.path.join(args.local_workspace, 'checkpoint_{:09d}.pth'.format(step)) |
| 46 | torch.save(model.state_dict(), CHECKPOINT_PATH) |
| 47 | print(colored('[MODEL]: ', 'yellow') + 'Saving the model at iteration {:d}: {:s}'.format( |
| 48 | step, os.path.basename(CHECKPOINT_PATH))) |
| 49 | elif epoch is not None: |
| 50 | CHECKPOINT_PATH = os.path.join(args.local_workspace, 'checkpoint_latest.pth') |
| 51 | torch.save({'model': model.state_dict(), |
| 52 | 'optimizer': optimizer.state_dict(), |
| 53 | 'epoch': epoch}, CHECKPOINT_PATH) |
| 54 | print(colored('[MODEL]: ', 'yellow') + 'Saving the model at epoch {:d}: {:s}'.format( |
| 55 | epoch, os.path.basename(CHECKPOINT_PATH))) |
| 56 | else: |
| 57 | raise "save model for a certain steps or epochs" |
| 58 | |
| 59 | |
| 60 | def load_checkpoint(args, model, optimizer): |