save model to checkpoint
(self, model, model_ema, contrast, optimizer, epoch)
| 105 | return start_epoch |
| 106 | |
| 107 | def save(self, model, model_ema, contrast, optimizer, epoch): |
| 108 | """save model to checkpoint""" |
| 109 | args = self.args |
| 110 | if args.local_rank == 0: |
| 111 | # saving the model to each instance |
| 112 | print('==> Saving...') |
| 113 | state = { |
| 114 | 'model': model.state_dict(), |
| 115 | 'contrast': contrast.state_dict(), |
| 116 | 'optimizer': optimizer.state_dict(), |
| 117 | 'epoch': epoch, |
| 118 | } |
| 119 | if isinstance(model_ema, torch.nn.Module): |
| 120 | state['model_ema'] = model_ema.state_dict() |
| 121 | if args.amp: |
| 122 | state['amp'] = amp.state_dict() |
| 123 | save_file = os.path.join(args.model_folder, 'current.pth') |
| 124 | torch.save(state, save_file) |
| 125 | if epoch % args.save_freq == 0: |
| 126 | save_file = os.path.join( |
| 127 | args.model_folder, 'ckpt_epoch_{}.pth'.format(epoch)) |
| 128 | torch.save(state, save_file) |
| 129 | # help release GPU memory |
| 130 | del state |
| 131 | |
| 132 | def train(self, epoch, train_loader, model, model_ema, contrast, |
| 133 | criterion, optimizer): |