预训练模型的参数保存与加载

1. 预训练参数保存

def savecheckpoint(model,dir):
    para=model.state_dict()
    torch.save(para, dir)

2. 预训练参数加载

def LoadParameter(_structure, _parameterDir):
    """

    :param _structure: model
    :param _parameterDir: 参数位置
    :return:
    """
    checkpoint = torch.load(_parameterDir)
    pretrained_state_dict = checkpoint['state_dict']
    model_state_dict = _structure.state_dict()

    for key in pretrained_state_dict:
        if key in model_state_dict:
            model_state_dict[key] = pretrained_state_dict[key]
    """
    another method:
    pretrained_dict = {k: v for k, v in pretrained_state_dict.items() if k in model_state_dict}
    model_state_dict.update(pretrained_dict)
    """

    _structure.load_state_dict(model_state_dict)
    model = torch.nn.DataParallel(_structure).cuda()

    return model