Load only model state while resuming training

Hi!

I have a use case where I've trained a model A with let's say 10 classes, and saved the state. Later, I want to load the same model and fine-tune it for 80 classes. To do this, I use the following function:

def resume_from_apex_checkpoint(fpath, model, optimizer=None, scheduler=None,
                                opt_level="O1", remove_classifier=False):
    """Resumes training from a checkpoint.

    This will load (1) model weights and (2) ``state_dict``
    of optimizer if ``optimizer`` is not None.

    Args:
        fpath (str): path to checkpoint.
        model (nn.Module): model.
        optimizer (Optimizer, optional): an Optimizer.
        scheduler (LRScheduler, optional): an LRScheduler.
        opt_level (str): optimization level if using apex.
        remove_classifier (bool): removes the classifier head while loading the
            model.

    Returns:
        int: start_epoch.

    Examples::
        >>> from utils import resume_from_apex_checkpoint
        >>> fpath = 'log/my_model/model.pth.tar-10'
        >>> start_epoch = resume_from_checkpoint(
        >>>     fpath, model, optimizer, scheduler
        >>> )
    """
    print('Loading APEX checkpoint from "{}"'.format(fpath))
    checkpoint = load_checkpoint(fpath)
    model, optimizer = amp.initialize(model.cuda(), optimizer,
                                      opt_level=opt_level, verbosity=0)
    model = nn.DataParallel(model)
    if remove_classifier:
        # this is a hack to load PLROSNet model trained on a different
        # number of classes. This will change if you use a different model.
        del checkpoint["state_dict"]["module.classifier1.weight"]
        del checkpoint["state_dict"]["module.classifier1.bias"]
        del checkpoint["state_dict"]["module.classifier2.weight"]
        del checkpoint["state_dict"]["module.classifier2.bias"]

    model.load_state_dict(checkpoint['state_dict'], strict=False)
    print('Loaded model weights')
    if optimizer is not None \
            and 'optimizer' in checkpoint.keys() \
            and not remove_classifier:
        optimizer.load_state_dict(checkpoint['optimizer'])
        print('Loaded optimizer')
    if scheduler is not None \
            and 'scheduler' in checkpoint.keys() \
            and not remove_classifier:
        scheduler.load_state_dict(checkpoint['scheduler'])
        print('Loaded scheduler')
    amp.load_state_dict(checkpoint['amp'])
    print("Initialized amp")
    start_epoch = checkpoint['epoch']
    print('Last epoch = {}'.format(start_epoch))
    if 'rank1' in checkpoint.keys():
        print('Last rank1 = {:.1%}'.format(checkpoint['rank1']))
    return start_epoch

For the fine-tuning task, I don't want to load the optimiser state, as I will be starting from scratch. AFAIK, apex stores the scaled values in the amp state so my understanding is it loads it for both the model, and the optimiser which might not be correct in my case. Please correct my understanding if I'm wrong and suggest a way to load it such that I load the state only for the model, and not the optimiser.

Start by tracing resume_from_apex_checkpoint, especially load_checkpoint, model.load_state_dict, optimizer.load_state_dict, scheduler.load_state_dict, and amp.load_state_dict. Check how checkpoint state is separated between model and optimizer during fine-tuning. Done means a documented or tested way to load model state without restoring optimizer or scheduler state.