pytorch在CPU和GPU上加载模型

最新推荐文章于 2024-08-08 23:29:37 发布

-牧野-

最新推荐文章于 2024-08-08 23:29:37 发布

阅读量2.7w

点赞数 19

分类专栏： pytorch

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

版权

pytorch 专栏收录该内容

4 篇文章 0 订阅

订阅专栏

pytorch允许把在GPU上训练的模型加载到CPU上，也允许把在CPU上训练的模型加载到GPU上。

CPU->CPU,GPU->GPU

torch.load('gen_500000.pkl')

GPU->CPU

torch.load('gen_500000.pkl', map_location=lambda storage, loc: storage)

或：

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))

单机多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 模式