Missing key(s) in state_dict | unexpected_keys | 模型加载 [torch.load] 报错【解决方案】

---

Missing key(s) in state_dict 报错分析

---

在执行 yolov3 pytorch 模型加载时遇到如下报错

---

类似报错之前其实也遇到过,今天总结一下

• 所运行代码

    # Initialize model
    model = Darknet(opt.cfg, imgsz)

    # Load weights
    attempt_download(weights)
    if weights.endswith('.pt'):  # pytorch format
        model.load_state_dict(torch.load(weights, map_location=device)['model'])

• 遇到的报错大致如下

RuntimeError: Error(s) in loading state_dict for Darknet:
        Missing key(s) in state_dict: "total_ops", "total_params", 

# 或者

RuntimeError: Error(s) in loading state_dict for DataParallel:
        Missing key(s) in state_dict: "module.total_ops", "module.total_params",

• 官方文档链接

该报错快速直接解决方法

model.load_state_dict 方法默认 strict=True

---

model.load_state_dict 时设置 strict=False 即可解决该报错

---

这种方式适合，模型本身没有什么问题，只是多了一些推理时不需要的 类似 "module_list.total_ops", "module_list.total_params", "module_list.0.total_ops", "module_list.0.total_params" 相关的统计变量

• 代码修改如下

model.load_state_dict(torch.load(weights, map_location=device)['model'], strict=False)

• python3 detect.py --source data/samples/ --device 0,1 --weights weights/yolov3-spp.pt 成功执行输出如下

Namespace(agnostic_nms=False, augment=False, cfg='cfg/yolov3-spp.cfg', classes=None, conf_thres=0.3, device='0,1', fourcc='mp4v', half=False, img_size=512, iou_thres=0.6, names='data/coco.names', output='output', save_txt=False, source='data/samples/', view_img=False, weights='weights/yolov3-spp.pt')

Using CUDA device0 _CudaDeviceProperties(name='Quadro RTX 5000', total_memory=16125MB)
           device1 _CudaDeviceProperties(name='Quadro RTX 5000', total_memory=16125MB)


  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
Model Summary: 225 layers, 6.29987e+07 parameters, 6.29987e+07 gradients, 118.0 GFLOPS
image 1/3 data/samples/bus.jpg: 512x384 4 persons, 1 buss, Done. (0.014s)
image 2/3 data/samples/dog.png: 384x512 1 bicycles, 1 trucks, 1 dogs, Done. (0.013s)
image 3/3 data/samples/zidane.jpg: 288x512 2 persons, 2 ties, Done. (0.012s)

Results saved to /home/zql/project/project21Next/yoloDir/yolov3/output
Done. (0.179s)

额外思路

提供这样一个思路，这种方法在该代码中，又遇到其它报错，但理论上它应该是可用的一种思路

相关代码

    # Initialize model
    model = Darknet(opt.cfg, imgsz)

    # Load weights
    attempt_download(weights)
    if weights.endswith('.pt'):  # pytorch format
        model.load_state_dict(torch.load(weights, map_location=device)['model'], strict=False)
        torch.save(model, 'weights/yolov3-spp-again.pt')
        # model.load_state_dict(torch.load(weights, map_location=device)['model'], strict=True)
    else:  # darknet format
        load_darknet_weights(model, weights)

---

【Unexpected key(s) in state_dict】- 解决思路

---

其它博客资料了解到的报错原因

• 在训练的时候，使用了多卡联合训练，导致 权重 keys 多了一个 前缀【之前有遇到这样的情况】

解决思路：遍历修改 keys ，去掉前缀 【更为通用】 | 或者在 缺少前缀的一方 添加 keys

---

这种情况，具体示例，下次遇到再做补充

去掉前缀 的代码思路如下

model_resnet101 = get_net()
model_resnet101.cuda()
model_resnet101.load_state_dict({k.replace('module.',''):v for k,v in torch.load("densenet169_rnn_fold_1_model_best_f1.pth.tar")['state_dict'].items()})

---