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