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

最新推荐文章于 2024-07-26 20:52:07 发布

原创最新推荐文章于 2024-07-26 20:52:07 发布 · 1.5k 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #预训练

深度学习同时被 3 个专栏收录

18 篇文章

订阅专栏

python

18 篇文章

订阅专栏

pytorch的使用

14 篇文章

订阅专栏

本文详细介绍了如何使用PyTorch保存和加载预训练模型的参数，包括定义保存函数和加载函数，确保模型能够在训练后被正确保存，并在需要时重新加载用于后续任务。

部署运行你感兴趣的模型镜像

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

您可能感兴趣的与本文相关的镜像