Function load_model_checkpoint

classification/main_accelerate.py:150–163

(config, model, accelerator)

def load_model_checkpoint(config, model, accelerator):
    if config.MODEL.RESUME:
        try:
            checkpoint = torch.load(config.MODEL.RESUME)['model']
            checkpoint = {k.replace('module.', ''): v for k, v in checkpoint.items()}
            model.load_state_dict(checkpoint)
        except:
            accelerator.load_state(config.MODEL.RESUME)
    elif config.MODEL.PRETRAINED:
        try:
            load_pretrained(config, model, logger)
        except:
            accelerator.load_state(config.MODEL.PRETRAINED)
    return model


def save_checkpoint(save_dir, accelerator, epoch, max_acc, config, lr_scheduler=None):