在对目标检测模型进行测试出现以下问题:
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for DataParallel:
Missing key(s) in state_dict: "module.vgg.0.weight", "module.vgg.0.bias"...........
Unexpected key(s) in state_dict: "vgg.0.weight", "vgg.0.bias", "vgg.5.weight"...................
RuntimeError: Error(s) in loading state_dict for DataParallel: size mismatch for module.loc.0.weight: copying a param with shape torch.Size([16, 512, 1, 1]) from checkpoint, the shape in current model is torch.Size([16, 512, 3, 3]). size mismatch for mo
问题原因:
1.模型训练时使用的GPU,nn.DataParallel
2.原模型和载入模型参数不匹配
解决代码如下:
# load net
num_classes = len(labelmap) + 1 # +1 for background
net = build_ssd('test', 300, num_classes).cuda() # initialize SSD
# net = nn.DataParallel(net).cuda()
state_dict = torch.load(args.trained_model)
new_state_dict = {}
for key, value in state_dict.items():
new_key = key
if not key.startswith('module.'):
new_key = 'module.' + key
new_state_dict[new_key] = value
# net.load_state_dict(new_state_dict,strict=False)
try:
net.load_state_dict(new_state_dict)
except RuntimeError as e:
print('Ignoring "' + str(e) + '"')
# checkpoint = torch.load(args.trained_model)
# # net.load_state_dict({k.replace('module.vgg',''):v for k,v in checkpoint['vgg'].items()})
# net.load_state_dict(checkpoint)
net.eval()
print('Finished loading model!')