| 261 | |
| 262 | |
| 263 | def save_checkpoint(config, |
| 264 | epoch, |
| 265 | model, |
| 266 | max_accuracy, |
| 267 | optimizer, |
| 268 | lr_scheduler, |
| 269 | scaler, |
| 270 | logger, |
| 271 | model_ema=None, |
| 272 | max_accuracy_ema=None, |
| 273 | ema_decay=None, |
| 274 | model_ems=None, |
| 275 | max_accuracy_ems=None, |
| 276 | ems_model_num=None, |
| 277 | best=None): |
| 278 | save_state = { |
| 279 | 'model': model.state_dict(), |
| 280 | 'optimizer': optimizer.state_dict(), |
| 281 | 'lr_scheduler': lr_scheduler.state_dict(), |
| 282 | 'max_accuracy': max_accuracy, |
| 283 | 'epoch': epoch, |
| 284 | 'config': config |
| 285 | } |
| 286 | if model_ema is not None: |
| 287 | save_state['model_ema'] = get_state_dict(model_ema) |
| 288 | if max_accuracy_ema is not None: |
| 289 | save_state['max_accuracy_ema'] = max_accuracy_ema |
| 290 | if ema_decay is not None: |
| 291 | save_state['ema_decay'] = ema_decay |
| 292 | if model_ems is not None: |
| 293 | save_state['model_ems'] = get_state_dict(model_ems) |
| 294 | if max_accuracy_ems is not None: |
| 295 | save_state['max_accuracy_ems'] = max_accuracy_ems |
| 296 | if ems_model_num is not None: |
| 297 | save_state['ems_model_num'] = ems_model_num |
| 298 | if config.AMP_OPT_LEVEL != 'O0': |
| 299 | # save_state['amp'] = amp.state_dict() |
| 300 | save_state['amp'] = scaler.state_dict() |
| 301 | if best is None: |
| 302 | save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth') |
| 303 | else: |
| 304 | save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{best}.pth') |
| 305 | logger.info(f'{save_path} saving......') |
| 306 | torch.save(save_state, save_path) |
| 307 | logger.info(f'{save_path} saved !!!') |
| 308 | |
| 309 | if dist.get_rank() == 0 and isinstance(epoch, int): |
| 310 | to_del = epoch - config.SAVE_CKPT_NUM * config.SAVE_FREQ |
| 311 | old_ckpt = os.path.join(config.OUTPUT, f'ckpt_epoch_{to_del}.pth') |
| 312 | if os.path.exists(old_ckpt): |
| 313 | os.remove(old_ckpt) |
| 314 | |
| 315 | |
| 316 | def get_grad_norm(parameters, norm_type=2): |