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

最新推荐文章于 2023-09-22 18:03:50 发布

Bug刺客

最新推荐文章于 2023-09-22 18:03:50 发布

阅读量1.2k

点赞数 2

分类专栏： PyTorch框架学习文章标签：多GPU模型

PyTorch框架学习专栏收录该内容

34 篇文章

订阅专栏

转自：https://blog.youkuaiyun.com/tsq292978891/article/details/83586935

背景

在公司用多卡训练模型，得到权值文件后保存，然后回到实验室，没有多卡的环境，用单卡训练，加载模型时出错，因为单卡机器上，没有使用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