赞
踩
参考链接:https://zhuanlan.zhihu.com/p/402293541
训练模型中使用到
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()})
用自己训练的分类网络替换目标检测中的主干网络,在训练目标检测网络的时候,使用了自己训练网络的预训练模型。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。