Function load_weights

utils.py:191–223

Load pretrained weights from given path.

(model: torch.nn.Module, pretrained_path: str)

def load_weights(model: torch.nn.Module, pretrained_path: str):
    """
    Load pretrained weights from given path.
    """
    if not pretrained_path:
        return model

    print(f'Loading weights from {pretrained_path}')
    state_dict = torch.load(pretrained_path)
    sd_keys = list(state_dict.keys())
    if 'state' in sd_keys:
        state_dict = state_dict['state']
        for k in list(state_dict.keys()):
            if k.startswith('module.'):
                state_dict["{}".format(k[len('module.'):])] = state_dict[k]
                del state_dict[k]

        model.load_state_dict(state_dict, strict=False)

    elif 'params' in sd_keys:
        state_dict = state_dict['params']
        for k in list(state_dict.keys()):
            if k.startswith('encoder.'):
                state_dict["{}".format(k[len('encoder.'):])] = state_dict[k]

            del state_dict[k]

        model.load_state_dict(state_dict, strict=True)
    else:
        model.load_state_dict(state_dict)

    print("Weights loaded successfully ")
    return model


def get_fs_labels(method: str, num_way: int, num_query: int, num_shot: int):