case1、GPU下DataParallel训练的模型----> CPU
问题:最近遇到一个问题,发下gpu训练的模型在cpu下不能跑,很难受。
结果发现原来是自己使用了DataParallel的问题,我保存的是整个模型。
model=DataParalle(model)
torch.save(‘xx.pkl’,model)
然后把模型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=‘cpu’)
model = model.module#才是你的模型
完整代码:
model = ModelArch(para).cuda(0) # 网络结构
model = torch.nn.DataParallel(model, device_ids=[0]) # 将model转为muit-gpus模式
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) # 载入weights
model.load_state_dict(checkpoint) # 用weights初始化网络
# 载入为单gpu模型
gpu_model = model.module # GPU-version
# 载入为cpu模型
model = ModelArch(para)
model.load_state_dict(gpu_model.state_dict())
torch.save(cpu_model.state_dict(), 'cpu_mode.pth') # cpu模型存储, 注意这里的state_dict后的()必须加上,否则报'function' object has no attribute 'copy'错误
还有另外一种方案:修改权重orderdic
中的名称
上述代码只有在模型在一个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)
case2、GPU ----> CPU
(1) Save on GPU, Load on CPU
save
torch.save(model.state_dict(), PATH)
load
device = torch.device("cpu")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
case3、GPU ----> GPU / CPU —> CPU
(2) Save on GPU, Load on GPU
save
torch.save(model.state_dict(), PATH)
load
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
case4、CPU ----> GPU
(3) Save on CPU, Load on GPU
save
torch.save(model.state_dict(), PATH)
load
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))
model.to(device)
case5、gpu 1 —> gpu 0
torch.load('modelparameters.pth', map_location={'cuda:1':'cuda:0'})