(self, file_root, epoch_idx, train_map, val_map, ema_start)
| 217 | return top1, loss, auc |
| 218 | |
| 219 | def save_checkpoint(self, file_root, epoch_idx, train_map, val_map, ema_start): |
| 220 | |
| 221 | file_name = os.path.join(file_root, |
| 222 | time.strftime('%Y%m%d-%H-%M', time.localtime()) + '-' + str(epoch_idx) + '.pth') |
| 223 | |
| 224 | if self.DDP: |
| 225 | stact_dict = self.netloc_.module.state_dict() |
| 226 | else: |
| 227 | stact_dict = self.netloc_.state_dict() |
| 228 | |
| 229 | torch.save( |
| 230 | { |
| 231 | 'epoch_idx': epoch_idx, |
| 232 | 'state_dict': stact_dict, |
| 233 | 'train_map': train_map, |
| 234 | 'val_map': val_map, |
| 235 | 'lr': self.lr_, |
| 236 | 'optimizer': self.optimizer_.state_dict(), |
| 237 | 'scheduler': self.scheduler_.state_dict() |
| 238 | }, file_name) |
| 239 | |
| 240 | if ema_start: |
| 241 | ema_file_name = os.path.join(file_root, |
| 242 | time.strftime('%Y%m%d-%H-%M', time.localtime()) + '-EMA-' + str(epoch_idx) + '.pth') |
| 243 | ema_stact_dict = self.ema_model.module.module.state_dict() |
| 244 | torch.save( |
| 245 | { |
| 246 | 'epoch_idx': epoch_idx, |
| 247 | 'state_dict': ema_stact_dict, |
| 248 | 'train_map': train_map, |
| 249 | 'val_map': val_map, |
| 250 | 'lr': self.lr_, |
| 251 | 'optimizer': self.optimizer_.state_dict(), |
| 252 | 'scheduler': self.scheduler_.state_dict() |
| 253 | }, ema_file_name) |
no outgoing calls
no test coverage detected