[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()})

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

Logo

技术共进,成长同行——讯飞AI开发者社区

更多推荐