(args, epoch, model, model_without_ddp, optimizer, loss_scaler)
| 293 | |
| 294 | |
| 295 | def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler): |
| 296 | output_dir = Path(args.output_dir) |
| 297 | epoch_name = str(epoch) |
| 298 | if loss_scaler is not None: |
| 299 | checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)] |
| 300 | for checkpoint_path in checkpoint_paths: |
| 301 | to_save = { |
| 302 | 'model': model_without_ddp.state_dict(), |
| 303 | 'optimizer': optimizer.state_dict(), |
| 304 | 'epoch': epoch, |
| 305 | 'scaler': loss_scaler.state_dict(), |
| 306 | 'args': args, |
| 307 | } |
| 308 | |
| 309 | save_on_master(to_save, checkpoint_path) |
| 310 | else: |
| 311 | client_state = {'epoch': epoch} |
| 312 | model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state) |
| 313 | |
| 314 | |
| 315 | def load_model(args, model_without_ddp, optimizer, loss_scaler): |
nothing calls this directly
no test coverage detected