raise RuntimeError(‘Error(s) in loading state_dict for {}:\n\t{}‘.format(

在对目标检测模型进行测试出现以下问题:

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!')

你可能感兴趣的:(pytorch,目标检测,python)