(self, save_name, save_path)
| 210 | return {'eval/loss': total_loss / total_num, 'eval/top-1-acc': top1, 'eval/top-5-acc': top5} |
| 211 | |
| 212 | def save_model(self, save_name, save_path): |
| 213 | if self.it < 1000000: |
| 214 | return |
| 215 | save_filename = os.path.join(save_path, save_name) |
| 216 | # copy EMA parameters to ema_model for saving with model as temp |
| 217 | self.model.eval() |
| 218 | self.ema.apply_shadow() |
| 219 | ema_model = deepcopy(self.model) |
| 220 | self.ema.restore() |
| 221 | self.model.train() |
| 222 | |
| 223 | torch.save({'model': self.model.state_dict(), |
| 224 | 'optimizer': self.optimizer.state_dict(), |
| 225 | 'scheduler': self.scheduler.state_dict(), |
| 226 | 'it': self.it + 1, |
| 227 | 'ema_model': ema_model.state_dict()}, |
| 228 | save_filename) |
| 229 | |
| 230 | self.print_fn(f"model saved: {save_filename}") |
| 231 | |
| 232 | def load_model(self, load_path): |
| 233 | checkpoint = torch.load(load_path) |
no test coverage detected