转自原文:https://blog.csdn.net/CV_YOU/article/details/86670188
有时候,我们用pytorch进行多卡GPUs训练时候,保存模型应该用下面语句:
torch.save(model.module.state_dict(), model_out_path)
但是忘记加module了,直接用
torch.save(model.state_dict(), model_out_path)
kwargs={'map_location':lambda storage, loc: storage.cuda(gpu_id)}
def load_GPUS(model,model_path,kwargs):
state_dict = torch.load(model_path,**kwargs)
# 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)
return model
单卡的模型加载代码如下:
kwargs={'map_location':lambda storage, loc: storage.cuda(gpu_id)}
def load_GPU(model,model_path,kwargs):
state_dict = torch.load(model_path,**kwargs)
# create new OrderedDict that does not contain `module.`
model.load_state_dict(state_dict)
return model