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

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

一、问题描述

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

       # save modelif num_gpu ==  1:torch.dule.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]=vnet.load_state_dict(new_state_dict)

参考链接：.html