当前位置:   article > 正文

[Error]pytorch加载模型权重出现问题Unexpected key(s) in state_dict: “module.backbone.0.weight“,

[Error]pytorch加载模型权重出现问题Unexpected key(s) in state_dict: “module.backbone.0.weight“,

[Error]pytorch加载模型权重出现问题Unexpected key(s) in state_dict: “module.backbone.0.weight”,

参考链接:https://zhuanlan.zhihu.com/p/402293541

问题原因:训练好的模型权重,加载进行测试时报错,原因是权重中含有的网络模型的各个layer与block的名称与网络不对应,且前缀加了个module.

原因分析:

训练模型中使用到

1 model = ResNet50()

2 model = nn.DataParallel(model).cuda()

加载预训练模型时使用

model = ResNet50().cuda()

checkpoint = torch.load(‘./ResNet50.pth’) # 预训练模型的地址

model.load_state_dict(checkpoint[‘state_dict’])

原因:由于用DataParallel训练的模型数据并行方式的,key中会包含”module“关键字

问题解决:

解决方法1:由于用DataParallel训练的模型数据并行方式的,key中会包含”module“关键字,加载时直接用:

model = ResNet50().cuda()

model = nn.DataParallel(model)

checkpoint = torch.load(‘./ResNet50.pth’) # 预训练模型的地址

model.load_state_dict(checkpoint[‘state_dict’])

解决方法2:

model = ResNet50().cuda()

checkpoint = torch.load(‘./ResNet50.pth’)

model.load_state_dict({k.replace(‘module.’,‘’):v for k,v in checkpoint[‘state_dict’].items()})

用自己训练的分类网络替换目标检测中的主干网络,在训练目标检测网络的时候,使用了自己训练网络的预训练模型。

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/Monodyee/article/detail/160087
推荐阅读
相关标签
  

闽ICP备14008679号