| 21 | |
| 22 | |
| 23 | def load_ema_checkpoint(config, model_ema, logger): |
| 24 | logger.info( |
| 25 | f'==============> Resuming form {config.MODEL.RESUME}....................' |
| 26 | ) |
| 27 | if config.MODEL.RESUME.startswith('https'): |
| 28 | checkpoint = torch.hub.load_state_dict_from_url(config.MODEL.RESUME, |
| 29 | map_location='cpu', |
| 30 | check_hash=True) |
| 31 | else: |
| 32 | checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu') |
| 33 | |
| 34 | assert isinstance(checkpoint, dict) |
| 35 | if 'model_ema' in checkpoint: |
| 36 | new_state_dict = OrderedDict() |
| 37 | for k, v in checkpoint['model_ema'].items(): |
| 38 | if model_ema.ema_has_module: |
| 39 | name = 'module.' + k if not k.startswith('module') else k |
| 40 | else: |
| 41 | name = k |
| 42 | new_state_dict[name] = v |
| 43 | msg = model_ema.ema.load_state_dict(new_state_dict, strict=False) |
| 44 | logger.info(msg) |
| 45 | logger.info('Loaded state_dict_ema') |
| 46 | else: |
| 47 | logger.warning( |
| 48 | 'Failed to find state_dict_ema, starting from loaded model weights' |
| 49 | ) |
| 50 | |
| 51 | max_accuracy_ema = 0 |
| 52 | if 'max_accuracy_ema' in checkpoint: |
| 53 | max_accuracy_ema = checkpoint['max_accuracy_ema'] |
| 54 | if 'ema_decay' in checkpoint: |
| 55 | model_ema.decay = checkpoint['ema_decay'] |
| 56 | return max_accuracy_ema |
| 57 | |
| 58 | |
| 59 | def load_checkpoint(config, model, optimizer, lr_scheduler, scaler, logger): |