pytorch加载多GPU模型和单GPU模型

有时候,我们用pytorch进行多卡GPUs训练时候,保存模型应该用下面语句:

torch.save(model.module.state_dict(), model_out_path)

但是忘记加module了,直接用

torch.save(model.state_dict(), model_out_path)

所以加载模型会遇到模型中参数名字多了module的关键字而报错,所以用下面加载模型的代码

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