KeyError: ‘model.2.cv2.bn‘

问题描述

在做YoloV5剪枝的时候，出现的问题。问题如下：

Traceback (most recent call last):
  File "D:\ObjectDetection\yolov5prune-main\prune.py", line 794, in <module>
    test_prune(opt.data,
  File "D:\ObjectDetection\yolov5prune-main\prune.py", line 452, in test_prune
    pruned_model = ModelPruned(maskbndict=maskbndict, cfg=pruned_yaml, ch=3).cuda()
  File "D:\ObjectDetection\yolov5prune-main\models\yolo.py", line 250, in __init__
    self.model, self.save, self.from_to_map = parse_pruned_model(self.maskbndict, deepcopy(self.yaml), ch=[ch])  # model, savelist
  File "D:\ObjectDetection\yolov5prune-main\models\yolo.py", line 653, in parse_pruned_model
    cv2out = int(maskbndict[named_m_cv2_bn].sum())
KeyError: 'model.2.cv2.bn'

分析原因

由于yoloV5有不同的版本，我们在训练的时候使用了预训练模型，但是预训练模型和这个版本的配置文件（yaml文件）不一致，所以在做剪枝的时候就会报错。

解决方法

在做稀疏训练的时候指定yaml模型文件。