pytorch多GPU训练和测试的问题： RuntimeError: Error(s) in loading state_dict for CPN:

本文链接：https://blog.youkuaiyun.com/yangzhengzheng95/article/details/88574200

本文探讨了在Windows10环境下，使用两块显卡进行PyTorch模型训练时，如何解决测试阶段导入模型遇到的state_dict错误。提供了三种解决方案，包括在测试时重新使用DataParallel，调整模型参数字典key，以及在保存模型时去除module前缀。

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

Windows10上，使用两块显卡，训练pytorch模型，语句如下：

model = network()
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)    # 多gpu上训练
model.cuda()

然后在测试时，导入模型会报如下类似错误：

RuntimeError: Error(s) in loading state_dict for CPN:
Missing key(s) in state_dict: "resnet.conv1.weight", "resnet.bn1.weight",
Unexpected key(s) in state_dict: "module.resnet.conv1.weight", "module.resnet.bn1.weight",

导入模型的语句：

model = network（）
model.load_state_dict(torch.load(model_state_path))
model.cuda()

报错原因分析如下：

在训练时，在多GPU上，我们使用了nn.DataParallel(model)对模型进行了包装，然后保存模型参数时，会在每个模型参数前添加module字段，作为模型预训练参数字典文件中的key, 因此为了保证在多GPU上训练的模型可以在测试时被正确导入，解决方法有两种：

方法一：

model = network()
model = nn.DataParallel(model)    # 添加了该句后，就能正常导入在多GPU上训练的模型参数了
model.load_state_dict(torch.load(model_state_path))
model.cuda()

方法二：

model = network.__dict__['CPN50']((192, 256), 1, pretrained=True)
model.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(model_state_path)['state_dict'].items()})
model.cuda()

方法三：

训练模型使用多个gpu，然后在测试时候不需要用单个gpu，因此在保存模型时应该吧module层去掉

if use_single_gpu:
    t.save(net.module.state_dict(), "model.pth")
else:
    t.save(net.state_dict(), "model.pth")