pytorch中 如何将gpu与gpu、gpu与cpu 在load时相互转化载入

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’ 如何解决？

您可能已经使用模型保存了模型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)

原文出处： https://www.ptorch.com/news/74.html

问题描述

有时在加载已训练好的模型时，会出现  out of memory  的错误提示，但仔细检测使用的GPU卡并没有再用且内存也没有超出。

经查阅发现原来是训练模型时使用的GPU卡和加载时使用的GPU卡不一样导致的。个人感觉，因为pytorch的模型中是会记录有GPU信息的，所以有时使用不同的GPU加载时会报错。 

解决方法

gpu之间的相互转换。即，将训练时的gpu卡转换为加载时的gpu卡。

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'})  # gpu 1 --> gpu 0

当然，你也可以直接将加载模型时使用的gpu卡改为和训练时的一样。但在多人使用一个服务器时，你想用的gpu卡已被使用，就需按上面方法转换gpu。

拓展

cpu --> cpu 或 gpu --> gpu

checkpoint = torch.load('modelparameters.pth')


model.load_state_dict(checkpoint)

cpu --> gpu 0

torch.load('modelparameters.pth', map_location=lambda storage, loc: storage.cuda(0))

gpu --> cpu

torch.load('modelparameters.pth', map_location=lambda storage, loc: storage)

gpu 1 --> gpu 0

torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'})

gpu训练的模型在cpu下不能跑，很难受。 
结果发现原来是自己使用了DataParallel的问题，我保存的是整个模型。

model=DataParalle(model)
torch.save('xx.pkl',model)
1
2
然后把模型copy到自己电脑上用cpu跑

model=torch.load('xx.pkl')
#此处报错了。。。。

原因是cpu环境不能直接导入gpu训练的DataParallel模型

所以换个策略,现在gpu把模型转化掉

model  = DataParallel(model)
....
real_model = model.module#这个才是你实际的模型，如果直接报错model的话，其实是保存了DataParallel(model)这个，这样会导致cpu环境下加载出错
torch.save(real_model,'xxx.pkl') #这样才是正确的保存模型方式，这样在cpu环境的模型才不会出错

其实还有一种方案，下面这种

model = torch.load('xxx.pkl', map_location=lambda storage, loc: storage)
model = model.module#才是你的模型

#上面那种貌似已经过时了,0.4.0的新方式
model = torch.load('xxx.pkl', map_location='cpu')
model = model.module#才是你的模型
==========

直接在GPU上加载：

 pretrain = torch.load(opt.pretrain_path)
 model.load_state_dict(pretrain['state_dict'])

将GPU模型加载在CPU上：

 pretrain = torch.load(opt.pretrain_path, map_location=lambda storage, loc: storage)
 from collections import OrderedDict
 new_state_dict = OrderedDict()
 for k, v in pretrain.items():
     if k=='state_dict':
         state_dict=OrderedDict()
         for keys in v:
             name = keys[7:]# remove `module.`
             state_dict[name] = v[keys]
         new_state_dict[k]=state_dict
     else:
         new_state_dict[k] = v
 model.load_state_dict(new_state_dict['state_dict'])