从官方文档中我们可以看到如下方法
torch.load('tensors.pt')
# 把所有的张量加载到CPU中
torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# 把所有的张量加载到GPU 1中
torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# 把张量从GPU 1 移动到 GPU 0
torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})
在cpu上加载预先训练好的GPU模型,有一种强制所有GPU张量在CPU中的方式:
torch.load('my_file.pt', map_location=lambda storage, loc: storage)
上述代码只有在模型在一个GPU
上训练时才起作用。如果我在多个GPU
上训练我的模型,保存它,然后尝试在CPU
上加载,我得到这个错误:KeyError: 'unexpected key "module.conv1.weight" in state_dict'
如何解决?
您可能已经使用模型保存了模型nn.DataParallel
,该模型将模型存储在该模型中module
,而现在您正试图加载模型DataParallel
。您可以nn.DataParallel
在网络中暂时添加一个加载目的,也可以加载权重文件,创建一个没有module
前缀的新的有序字典,然后加载它。
# original saved file with DataParallel
state_dict = torch.load('myfile.pth.tar')
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)