(args, epoch, model, model_without_ddp, optimizer, loss_scaler, ema_params=None, epoch_name=None)
| 304 | |
| 305 | |
| 306 | def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler, ema_params=None, epoch_name=None): |
| 307 | if epoch_name is None: |
| 308 | epoch_name = str(epoch) |
| 309 | output_dir = Path(args.output_dir) |
| 310 | checkpoint_path = output_dir / ('checkpoint-%s.pth' % epoch_name) |
| 311 | |
| 312 | # ema |
| 313 | if ema_params is not None: |
| 314 | ema_state_dict = copy.deepcopy(model_without_ddp.state_dict()) |
| 315 | for i, (name, _value) in enumerate(model_without_ddp.named_parameters()): |
| 316 | assert name in ema_state_dict |
| 317 | ema_state_dict[name] = ema_params[i] |
| 318 | else: |
| 319 | ema_state_dict = None |
| 320 | |
| 321 | to_save = { |
| 322 | 'model': model_without_ddp.state_dict(), |
| 323 | 'model_ema': ema_state_dict, |
| 324 | 'optimizer': optimizer.state_dict(), |
| 325 | 'epoch': epoch, |
| 326 | 'scaler': loss_scaler.state_dict(), |
| 327 | 'args': args, |
| 328 | } |
| 329 | save_on_master(to_save, checkpoint_path) |
| 330 | |
| 331 | |
| 332 | def all_reduce_mean(x): |
nothing calls this directly
no test coverage detected