(accelerator, config, lr_scheduler=None)
| 178 | |
| 179 | |
| 180 | def load_checkpoint_if_needed(accelerator, config, lr_scheduler=None): |
| 181 | setup_autoresume(config) |
| 182 | save_dir = config.MODEL.RESUME |
| 183 | if not save_dir: |
| 184 | return 0.0 |
| 185 | accelerator.load_state(save_dir) |
| 186 | checkpoint = torch.load(os.path.join(save_dir, 'additional_state.pth'), map_location='cpu') |
| 187 | if lr_scheduler is not None: |
| 188 | logger.info('resuming lr_scheduler') |
| 189 | lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) |
| 190 | config.defrost() |
| 191 | config.TRAIN.START_EPOCH = checkpoint['epoch'] + 1 |
| 192 | config.freeze() |
| 193 | max_acc = checkpoint.get('max_acc', 0.0) |
| 194 | logger.info(f"=> loaded successfully {config.MODEL.RESUME} (epoch {checkpoint['epoch']})") |
| 195 | return max_acc |
| 196 | |
| 197 | |
| 198 | def log_model_statistic(model_wo_ddp): |
no test coverage detected