pytorch多GPU训练保存的模型，在单GPU环境下加载出错

最新推荐文章于 2025-05-19 21:11:21 发布

tang-0203

最新推荐文章于 2025-05-19 21:11:21 发布

阅读量9.1k

点赞数 3

CC 4.0 BY-SA版权

分类专栏： Pytorch学习文章标签：多GPU 模型加载

本文链接：https://blog.youkuaiyun.com/tsq292978891/article/details/83586935

Pytorch学习专栏收录该内容

24 篇文章

订阅专栏

在没有多卡的环境中，尝试加载由多卡训练的模型会导致错误，因为权值参数带有module前缀。本文介绍了两种解决方案：一是用DataParallel包装单卡模型；二是自定义Load函数，适配不同环境。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

背景
在公司用多卡训练模型，得到权值文件后保存，然后回到实验室，没有多卡的环境，用单卡训练，加载模型时出错，因为单卡机器上，没有使用DataParallel来加载模型，所以会出现加载错误。
原因
DataParallel包装的模型在保存时，权值参数前面会带有module字符，然而自己在单卡环境下，没有用DataParallel包装的模型权值参数不带module。本质上保存的权值文件是一个有序字典。
解决方法
1.在单卡环境下，用DataParallel包装模型。

2.自己重写Load函数，灵活。

from collections import OrderedDict
def myOwnLoad(model, check):
    modelState = model.state_dict()
    tempState = OrderedDict()
    for i in range(len(check.keys())-2):
        print modelState.keys()[i], check.keys()[i]
        tempState[modelState.keys()[i]] = check[check.keys()[i]]
    temp = [[0.02]*1024 for i in range(200)]  # mean=0, std=0.02
    tempState['myFc.weight'] = torch.normal(mean=0, std=torch.FloatTensor(temp)).cuda()
    tempState['myFc.bias']   = torch.normal(mean=0, std=torch.FloatTensor([0]*200)).cuda()

    model.load_state_dict(tempState)
    return model