多GPU训练，加载模型时报错

星月夜语

于 2021-04-02 08:31:53 发布

阅读量2.1k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： python 深度学习基础知识文章标签：多GPU 加载报错

本文链接：https://blog.csdn.net/ljh618625/article/details/115389817

深度学习基础知识同时被 2 个专栏收录

40 篇文章 ¥49.90 ¥99.00

订阅专栏

超级会员免费看

python

34 篇文章

订阅专栏

在使用多个GPU进行训练并保存模型时，未添加'module'导致加载时出现错误。解决方法是在加载模型时删除多余的'module'，通过打印模型路径来定位问题。参考链接提供了解决方案。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

一、问题描述

多个GPU 训练，保存时没有加module , 导致加载模型时报错。正确写法应该如下：

       # save model
        if num_gpu ==  1:
            torch.save(model.module.state_dict(), os.path.join(opt.outf, 'net.pth'))
        else:
            torch.save(model.state_dict(), os.path.join(opt.outf, 'net.pth'))

二、解决方法

load 模型时，删除多余的module. 可以打印下面代码中的pth , 查看问题症结

具体代码如下：

        print("load pre_training weight. ")
        pth = torch.load(os.path.join(opt.outf, 'net.pth'))
        new_state_dict = OrderedDict()
        for k, v in pth.items():
            name =  k[7:] # remove  'module'
            new_state_dict[name]=v
        net.load_state_dict(new_state_dict)

参考链接：http://www.likecs.com/show-99262.html