yolov5剪枝,支持v2,v3,v4版本的yolov5

Overview

yolov5_prune

本项目基于tanluren/yolov3-channel-and-layer-pruning实现,将项目扩展到yolov5上。

项目的基本流程是,使用ultralytics/yolov5训练自己的数据集,在模型性能达到要求但速度未达到要求时,对模型进行剪枝。首先是稀疏化训练,稀疏化训练很重要,如果模型稀疏度不够,剪枝比例过大会导致剪枝后的模型map接近0。剪枝完成后对模型进行微调回复精度。

本项目使用的yolov5为第四版本。 yolov5第三版本参考yolov5-v3-prune yolov5第二版本参考yolov5-v2-prune

TODO: 增加m,l,x的模型剪枝,如果有时间的话。>-<

PS:在开源数据集和不能开源的数据集上模型均剪枝成功。

实例流程

数据集下载dataset

STEP1:基础训练

附件:训练记录

STEP2:稀疏训练

附件:稀疏训练记录

STEP3:八倍通道剪枝

附件:剪枝后模型

STEP4:微调finetune

附件:微调训练记录

STEP4:微调finetune,使用蒸馏技术优化模型,效果由于单纯的微调模型

附件:微调蒸馏训练记录

剪枝步骤

STEP1:基础训练

yolov5
示例代码

python train.py --img 640 --batch 8 --epochs 50 --weights weights/yolov5s_v4.pt --data data/coco_hand.yaml --cfg models/yolov5s.yaml --name s_hand

STEP2:稀疏训练

--prune 0 适用于通道剪枝策略一,--prune 1 适用于其他剪枝策略。
yolov5
示例代码

python train_sparsity.py --img 640 --batch 8 --epochs 50 --data data/coco_hand.yaml --cfg models/yolov5s.yaml --weights runs/train/s_hand/weights/last.pt --name s_hand_sparsity -sr --s 0.001 --prune 1

STEP3:通道剪枝策略一

不对shortcut直连的层进行剪枝,避免维度处理。

python prune_yolov5s.py --cfg cfg/yolov5s.cfg --data data/fangweisui.data --weights weights/yolov5s_prune0.pt --percent 0.8

STEP3:通道剪枝策略二

对shortcut层也进行了剪枝,剪枝采用每组shortcut中第一个卷积层的mask。

python shortcut_prune_yolov5s.py --cfg cfg/yolov5s.cfg --data data/fangweisui.data --weights weights/yolov5s_prune1.pt --percent 0.3

STEP3:通道剪枝策略三

先以全局阈值找出各卷积层的mask,然后对于每组shortcut,它将相连的各卷积层的剪枝mask取并集,用merge后的mask进行剪枝。

python slim_prune_yolov5s.py --cfg cfg/yolov5s.cfg --data data/fangweisui.data --weights weights/yolov5s_prune1.pt --global_percent 0.8 --layer_keep 0.01

STEP3:八倍通道剪枝

在硬件部署上发现,模型剪枝率相同时,通道数为8的倍数速度最快。(采坑:需要将硬件性能开启到最大)
示例代码

python slim_prune_yolov5s_8x.py --cfg cfg/yolov5s_v4_hand.cfg --data data/oxfordhand.data --weights weights/last_v4s.pt --global_percent 0.5 --layer_keep 0.01 --img_size 640

STEP4:微调finetune

yolov5
示例代码

python prune_finetune.py --img 640 --batch 8 --epochs 50 --data data/coco_hand.yaml --cfg ./cfg/prune_0.5_keep_0.01_8x_yolov5s_v4_hand.cfg --weights ./weights/prune_0.5_keep_0.01_8x_last_v4s.pt --name s_hand_finetune

STEP4:微调finetune,使用蒸馏技术优化模型

yolov5
示例代码

python prune_finetune.py --img 640 --batch 8 --epochs 50 --data data/coco_hand.yaml --cfg ./cfg/prune_0.5_keep_0.01_8x_yolov5s_v4_hand.cfg --weights ./weights/prune_0.5_keep_0.01_8x_last_v4s.pt --name s_hand_finetune_distill --distill

STEP5:剪枝后模型推理

yolov5
示例代码

python prune_detect.py --weights weights/last_s_hand_finetune.pt --img  640 --conf 0.7 --save-txt --source inference/images
Issues
  • map before prune?

    map before prune?

    command: python prune_yolov5s.py --cfg cfg/yolov5s.cfg --data data/fangweisui.data --weights weights/yolov5s_prune0.pt --percent 0.8 code used: [https://github.com/ZJU-lishuang/yolov5]

    question: +------------+----------+----------+ | Metric | Before | After | +------------+----------+----------+ | mAP | 0.025648 | 0.000000 | | Parameters | 7263185 | 1790052 | | Inference | 0.0165 | 0.0172 | +------------+----------+----------+ so when eval the ori model,what is the format of annos ?like normalized [classid,x_center,y_center,w,h] or coco format?

    opened by codylcs 34
  • 如何参考yolov5s.cfg得到yolov5x.cfg

    如何参考yolov5s.cfg得到yolov5x.cfg

    @ZJU-lishuang 你好,yolov5s.yaml和yolov5x.yaml只有width_multiple和depth_multiple两个参数不一样 如果要得到yolov5x.cfg对yolov5x进行剪枝,在yolov5s.cfg的基础上如何修改可以得到yolov5x.cfg? 刚接触这方面的大小不知道该如何入手

    opened by yao-zheng-yi 19
  • 'Conv2d' object has no attribute 'conv'

    'Conv2d' object has no attribute 'conv'

    你好 我尝试使用自己的数据集训练模型,按照readme的流程,在使用8倍通道剪枝的时候出现了问题。 Traceback (most recent call last): File "slim_prune_yolov5s_8x.py", line 342, in copy_weight_v4(modelyolov5, model) File "slim_prune_yolov5s_8x.py", line 23, in copy_weight_v4 copy_conv(cspnet1.cv2, model.module_list[3]) File "slim_prune_yolov5s_8x.py", line 13, in copy_conv conv_dst[0] = conv_src.conv File "/home/cwy/.conda/envs/yolo/lib/python3.8/site-packages/torch/nn/modules/module.py", line 778, in getattr raise ModuleAttributeError("'{}' object has no attribute '{}'".format( torch.nn.modules.module.ModuleAttributeError: 'Conv2d' object has no attribute 'conv'

    其中,cfg文件我是复制yolov5s_hand.cfg,修改了class,filters和anchors三个参数 data文件的内容如下: classes= 10 train = /data/cwy/CardNumber/yolov5_random/train.txt valid = /data/cwy/CardNumber/yolov5_random/test.txt names = data/CardNumber.names backup = /data/cwy/CardNumber/yolov5_random/

    改了很久也没能改对,请问这个问题该如何解决

    opened by yao-zheng-yi 14
  • 为什么稀疏训练后的权重文件有54.1M

    为什么稀疏训练后的权重文件有54.1M

    python train_sparsity.py --img 640 --batch 16 --epochs 50 --data data/coco_hand.yaml --cfg models/yolov5s.yaml --weights runs/train/s_hand/weights/last.pt --name s_hand_sparsity -sr --s 0.001 --prune 0

    而且这个--s参数没有啊,大佬帮忙看看什么问题

    opened by ztrong-forever 12
  • python -c

    python -c "from modelsori import *; convert('cfg/sar-yolov4-csp.cfg', 'weights/sar2-best.pt')"

    loaded weights from weights/converted.weights

    let's test the original model first: Class Images Targets P R mAP F1: 0%| | 0/174 [00:00<?, ?it/s] 转换出来的.weights不能测试,这是为什么

    opened by lccui 11
  • 减枝后的model无法在yolov5上运行

    减枝后的model无法在yolov5上运行

    你好,我用自己的数据在这个项目里剪枝后,打算在原yolov5的项目中test,但是报错; 然后我用你在tutorial 里给的8倍通道剪枝的model去跑,结果还是一样。报错如下: Traceback (most recent call last): File "test.py", line 309, in test(opt.data, File "test.py", line 55, in test model = attempt_load(weights, map_location=device) # load FP32 model File "/home/bai/Desktop/yolov5-v3-test/models/experimental.py", line 118, in attempt_load model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model AttributeError: 'collections.OrderedDict' object has no attribute 'float' 请问是model 需要再转换一下格式么? 但是我看它已经是.pt格式了,实在不知道咋整。

    opened by maxiaotian520 10
  • 请问微调训练后,怎么测试得到每一类的p、r、map值?

    请问微调训练后,怎么测试得到每一类的p、r、map值?

    作者您好,我正常的进行了训练、稀疏训练、剪枝(采用了prune0),最后微调恢复训练。 我想测试微调后的模型性能,使用了您yolov5-v4库里面的test.py,但是出现了以下的报错。

    PS F:\PycharmProjects\pythonProject4\yolov5-v4> python test.py --weights test-models/prune0.7/weights/prune_0.7_sewer-yolov5s-sparsity0-last_finetune-best.pt --data data/coco128.yaml --batch-size 1 --img-size 640 --iou-thres 0.5 --task
     val --verbose --name prune_0.7_sewer-yolov5s-sparsity0-last_finetune-best-val
    Namespace(augment=False, batch_size=1, conf_thres=0.001, data='data/coco128.yaml', device='', exist_ok=False, img_size=640, iou_thres=0.5, name='prune_0.7_sewer-yolov5s-sparsity0-last_finetune-best-val', project='runs/test', save_conf=
    False, save_hybrid=False, save_json=False, save_txt=False, single_cls=False, task='val', verbose=True, weights=['test-models/prune0.7/weights/prune_0.7_sewer-yolov5s-sparsity0-last_finetune-best.pt'])
    Using torch 1.7.1 CUDA:0 (Quadro P5000, 16384.0MB)
    
    
    Traceback (most recent call last):
      File "test.py", line 566, in <module>
        test(opt.data,
      File "test.py", line 315, in test
        imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size
      File "D:\Anaconda3\envs\yolo_env\lib\site-packages\torch\nn\modules\module.py", line 778, in __getattr__
        raise ModuleAttributeError("'{}' object has no attribute '{}'".format(
    torch.nn.modules.module.ModuleAttributeError: 'Darknet' object has no attribute 'stride'
    

    然后我看回来您yolov5_prune库里面也有一个test.py,于是乎我也试着用一下看能不能得到微调后的模型数据,结果也报错了。(而且没有--task val/test,--verbose也没有)

    PS F:\PycharmProjects\pythonProject4\yolov5-v4-prune> python test.py --cfg test-models/prune0.7/cfg/prune_0.7_yolov5s_v4.cfg --weights test-models/prune0.7/weights/prune_0.7_sewer-yolov5s-sparsity0-last_finetune-best.pt --data data/coc
    o_128img.data --batch-size 1 --img-size 640 --iou-thres 0.5
    Namespace(batch_size=1, cfg='test-models/prune0.7/cfg/prune_0.7_yolov5s_v4.cfg', conf_thres=0.001, data='data/coco_128img.data', device='', img_size=640, iou_thres=0.5, nms_thres=0.5, save_json=False, weights='test-models/prune0.7/weig
    hts/prune_0.7_sewer-yolov5s-sparsity0-last_finetune-best.pt')
    Using CUDA device0 _CudaDeviceProperties(name='Quadro P5000', total_memory=16384MB)
    
    Traceback (most recent call last):
      File "test.py", line 324, in <module>
        test(opt.cfg,
      File "test.py", line 140, in test
        model.load_state_dict(torch.load(weights, map_location=device)['model'])
      File "D:\Anaconda3\envs\yolo_env\lib\site-packages\torch\nn\modules\module.py", line 1025, in load_state_dict
        state_dict = state_dict.copy()
      File "D:\Anaconda3\envs\yolo_env\lib\site-packages\torch\nn\modules\module.py", line 778, in __getattr__
        raise ModuleAttributeError("'{}' object has no attribute '{}'".format(
    torch.nn.modules.module.ModuleAttributeError: 'Darknet' object has no attribute 'copy'
    

    请问作者您有一份可以测试剪枝后和微调后模型的代码?感谢。

    opened by zxsitu 9
  • 在剪枝的时候报错

    在剪枝的时候报错

    你好,我基于u版的yolo v5s训练出了自己的模型,用你的代码进行了稀疏训练。在使用剪枝策略1时代码报错了,信息如下:

    (pytorch) D:\CNN\yolov5\zju_prune\yolov5_prune-4>python prune_yolov5s.py --cfg D:\CNN\yolov5\yolov5\models\yolov5s.yaml --data D:\CNN\yolov5\zju_prune\yolov5_prune-4\datasets_3_aug_yolo\data.yaml --weights D:\CNN\yolov5\zju_prune\yolov5-v4-main\runs\train\s_sparsity\weights\best.pt --percent 0.8 Namespace(cfg='D:\\CNN\\yolov5\\yolov5\\models\\yolov5s.yaml', data='D:\\CNN\\yolov5\\zju_prune\\yolov5_prune-4\\datasets_3_aug_yolo\\data.yaml', img_size=416, percent=0.8, weights='D:\\CNN\\yolov5\\zju_prune\\yolov5-v4-main\\runs\\train\\s_sparsity\\weights\\best.pt') Traceback (most recent call last): File "prune_yolov5s.py", line 338, in <module> model = Darknet(opt.cfg, (img_size, img_size)).to(device) File "D:\CNN\yolov5\zju_prune\yolov5_prune-4\modelsori.py", line 254, in __init__ self.module_defs = parse_model_cfg(cfg) File "D:\CNN\yolov5\zju_prune\yolov5_prune-4\utils\parse_config.py", line 18, in parse_model_cfg key, val = line.split("=") ValueError: not enough values to unpack (expected 2, got 1)

    我看你使用了data data/fangweisui.data这个文件作为训练数据设置,我不知道你的.data文件格式是怎样的。我使用的是yaml文件,会不会是这里有影响呢?

    opened by dada1437903138 7
  • 稀疏训练后的通道剪枝策略得到的map都是0

    稀疏训练后的通道剪枝策略得到的map都是0

    楼主您好,我使用的是yolov5-6网络。 (更新:发现剪枝率设置为0.1后。各项指标依旧下降为0) 因为在您的yolov5_prune工程中没有找到稀疏训练的train_sparsity代码,所以使用的是您仓库下的yolov5-v6工程里的train_sparsity代码,训练数据集也是您提供的coco_hand。 训练步骤:(不知道是不是这么操作) 1、我稀疏训练了200epoch,在100epoch时因为map没有继续上涨程序自动停止了。 2、得到的权重best.pt用于yolov5_prune工程中的prune_yolov5s.py进行剪枝。模型cfg文件使用的就是您代码中提供的yolov5s_v6_hand.cfg。 出现问题: prune_yolov5s.py在运行的时候,test的current model的P R Map和我train_sparsity的值不一样,差挺多的 且,剪枝后P R mAP F1全部为0。 运行截图 稀疏训练 图片 剪枝 图片 图片

    opened by dddmmmyyy1998 6
  • infer time and map (yolov5l prune vs  yolov5s)?

    infer time and map (yolov5l prune vs yolov5s)?

    在精度和耗时方面yolov5做了一个权衡,https://github.com/ultralytics/yolov5/issues/304,这里Glenn Jocher并不是特别建议自己作剪枝,如果追求更改的精度,可以选用大模型,如果追求更快的速度,用yolov5_s,如果用大模型去做剪枝和yolov5_s作对比的话,速度和精度会有优势吗?有相关实验结论吗?(yolov3元模型比较大的,200M以上,才有了剪枝的必要,yolov5同样适用剪枝吗?另一个问题还没有关闭,https://github.com/ZJU-lishuang/yolov5_prune/issues/4,发现用yolov5_v2得到的模型还没有直接用yolov5s_v3得到的模型精度高,所以有此疑惑,望不嫌麻烦,给点建议,不胜感激)

    opened by codylcs 6
  • 你好,我在 STEP3:通道剪枝 遇上问题,想请你帮忙看看

    你好,我在 STEP3:通道剪枝 遇上问题,想请你帮忙看看

    python prune_yolov5s.py --cfg cfg/yolov5s_v4.cfg --data data/myvoc.data --weights best.pt --percent 0.8 后出现以下错误: Traceback (most recent call last): File "prune_yolov5s.py", line 361, in origin_model_metric = eval_model(model) File "prune_yolov5s.py", line 348, in eval_model = lambda model:test(opt.cfg, opt.data, File "D:\Gitee\yolov5_prune-4\test.py", line 181, in test inf_out, train_out = model(imgs) # inference and training outputs File "C:\Users\TXY\anaconda3\envs\yolo_v5\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl result = self.forward(*input, **kwargs) File "D:\Gitee\yolov5_prune-4\modelsori.py", line 293, in forward x = module(x, img_size) File "C:\Users\TXY\anaconda3\envs\yolo_v5\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl result = self.forward(*input, **kwargs) File "D:\Gitee\yolov5_prune-4\modelsori.py", line 178, in forward p = p.view(bs, self.na, self.nc + 5, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous() # prediction RuntimeError: shape '[16, 3, 85, 28, 52]' is invalid for input of size 489216 你是否遇到过该错误呢? 还有请问通道剪枝命令命令 python prune_yolov5s.py --cfg cfg/yolov5s.cfg --data data/fangweisui.data --weights weights/yolov5s_prune0.pt --percent 0.8 中,yolov5s.cfg和fangweisui.data文件是否做了修改呢?是否可以上传一下呢?

    opened by Moon404 5
  • 为什么剪枝之后在nx开发板上运行速度也没有提升

    为什么剪枝之后在nx开发板上运行速度也没有提升

    剪枝后的onnx通过python版tensorrt转成engine文件,但跑在nx上的速度和PT生成engine速度一样,还有疑惑的是剪枝后生成的模型居然比PT文件直接转engine文件大不少。 我的流程是先原yolov5v4代码训练,然后换成楼主提供的v6代码做稀疏训练,接着将模型放到prune代码里进行剪枝,生成的cfg和pt再进行微调训练,转成onnx,最后用py tensorrt生成engine,c++推理,如果流程有什么问题恳请大佬请教

    opened by WEI-6 0