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 模式