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

最新推荐文章于 2024-07-26 20:52:07 发布

原创最新推荐文章于 2024-07-26 20:52:07 发布 · 1w 阅读

16 ·

CC 4.0 BY-SA版权

深度学习专栏收录该内容

39 篇文章

订阅专栏

本文详细解析了使用PyTorch框架在多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

您可能感兴趣的与本文相关的镜像