pytorch模型与数据在GPU和CPU上移动

阿委困的不能行

已于 2022-08-09 16:32:55 修改

阅读量2.6k

点赞数

文章标签： pytorch 深度学习 python

于 2020-08-31 09:51:18 首次发布

原文链接：https://blog.csdn.net/dcrmg/article/details/79503978

版权

原文链接：https://blog.csdn.net/dcrmg/article/details/79503978
pytorch允许把在GPU上训练的模型加载到CPU上，也允许把在CPU上训练的模型加载到GPU上。

CPU->CPU,GPU->GPU

torch.load('gen_500000.pkl')

GPU->CPU

torch.load('xxx.pkl', map_location='cpu')  # pytorch0.4.0及以上版本

CPU->GPU1

torch.load('gen_500000.pkl', map_location=lambda storage, loc: storage.cuda(1))
#或者
torch.load('gen_500000.pkl', map_location={'cpu':'cuda:1'})

单机多GPU上训练的模型加载
在gpu或cpu上加载一个在多gpu上训练的模型时，会报类似如下错误：

Unexpected key(s) in state_dict: "module.pretrained.conv1.0.weight", "module.pretrained.conv1.1.weight"

这是因为在多gou上训练的模型在保存时候在参数名前多加了一个“module.”前缀，加载的时候把这个前缀去掉就行了：

model = get_model(args.model)
device = torch.device('cpu')
# device = torch.device('cuda:0')
state_dict = torch.load('xxx.pth', map_location=device)
 
from collections import OrderedDict
state_dict_new = OrderedDict()
for k, v in state_dict.items():
    name = k[7:]  # 去掉 `module.`
    state_dict_new[name] = v
 
model.load_state_dict(state_dict_new)
 
model.eval() #eval 模式