pytorch中 如何将gpu与gpu、gpu与cpu 在load时相互转化载入
有时候我们在CPU上训练的模型,因为一些原因,切换到GPU上,或者在GPU上训练的模型,因为条件限制,切换到CPU上。 GPU上训练模型时,将权重加载到CPU的最佳方式是什么?今天我们来讨论一下:
提取模型到指定的设备
从官方文档中我们可以看到如下方法
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)
Q:上述代码只有在模型在一个GPU上训练时才起作用。如果我在多个GPU上训练我的模型,保存它,然后尝试在CPU上加载,我得到这个错误:KeyError: ‘unexpected key “module.conv1.weight” in state_dict’