| 57 | |
| 58 | |
| 59 | def load_checkpoint(config, model, optimizer, lr_scheduler, scaler, logger): |
| 60 | logger.info( |
| 61 | f'==============> Resuming form {config.MODEL.RESUME}....................' |
| 62 | ) |
| 63 | if config.MODEL.RESUME.startswith('https'): |
| 64 | checkpoint = torch.hub.load_state_dict_from_url(config.MODEL.RESUME, |
| 65 | map_location='cpu', |
| 66 | check_hash=True) |
| 67 | else: |
| 68 | checkpoint = torch.load(config.MODEL.RESUME, map_location='cpu') |
| 69 | |
| 70 | print('resuming model') |
| 71 | |
| 72 | model_checkpoint = checkpoint['model'] |
| 73 | msg = model.load_state_dict(model_checkpoint, strict=False) |
| 74 | logger.info(msg) |
| 75 | max_accuracy = 0.0 |
| 76 | if not config.EVAL_MODE and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint: |
| 77 | if optimizer is not None: |
| 78 | print('resuming optimizer') |
| 79 | try: |
| 80 | optimizer.load_state_dict(checkpoint['optimizer']) |
| 81 | except: |
| 82 | print('resume optimizer failed') |
| 83 | if lr_scheduler is not None: |
| 84 | print('resuming lr_scheduler') |
| 85 | lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) |
| 86 | config.defrost() |
| 87 | config.TRAIN.START_EPOCH = checkpoint['epoch'] + 1 |
| 88 | config.freeze() |
| 89 | if 'amp' in checkpoint and config.AMP_OPT_LEVEL != 'O0' and checkpoint['config'].AMP_OPT_LEVEL != 'O0': |
| 90 | scaler.load_state_dict(checkpoint['amp']) |
| 91 | logger.info( |
| 92 | f"=> loaded successfully {config.MODEL.RESUME} (epoch {checkpoint['epoch']})" |
| 93 | ) |
| 94 | if 'max_accuracy' in checkpoint: |
| 95 | max_accuracy = checkpoint['max_accuracy'] |
| 96 | |
| 97 | del checkpoint |
| 98 | torch.cuda.empty_cache() |
| 99 | |
| 100 | return max_accuracy |
| 101 | |
| 102 | |
| 103 | def load_pretrained(config, model, logger): |