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