YOLOv5s 剪枝

yolov5_prune

yolov5_prune源码地址：https://github.com/ZJU-lishuang/yolov5_prune
yolov5-v4源码地址：https://github.com/ZJU-lishuang/yolov5-v4

环境配置

python = 3.7
torch = 1.7.0+cu101
torchvision = 0.8.1+cu101

剪枝步骤

基础训练

YOLOv5自定义数据集训练：https://zhouchen.blog.csdn.net/article/details/108122600

在自己使用yolov5训练数据集时，因yolov5的warmup和余弦退火机制

应尽量将epochs数量调大，训练效果才能较优

在yolov5-v4项目运行

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

yolov5s.pt下载地址：
https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt
coco_hand.yaml：根据数据集自己编写
yolov5s.yaml：yolov5-v4项目下

稀疏训练

在yolov5-v4项目运行
–prune 0 适用于通道剪枝策略一，–prune 1 适用于其他剪枝策略。

python train_sparsity.py --img 640 --batch 2 --epochs 120 --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

last.pt：基础训练完后生成的.pt文件

通道剪枝

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

在yolov5_prune项目运行

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

修改yolov5s.cfg文件


coco-hand.data

coco_hand.names

微调finetune

在yolov5-v4项目运行

python prune_finetune.py --img 640 --batch 2 --epochs 120 --data data/coco_hand.yaml --cfg ./cfg/prune_0.8_yolov5s_v4.cfg --weights ./weights/prune_0.8_last.pt --name s_hand_finetune

剪枝后模型推理

在yolov5-v4项目运行

python prune_detect.py --weights weights/last_s_hand_finetune.pt --img  640 --conf 0.7 --save-txt --source inference/images

last_s_hand_finetune.pt:经过微调后训练获得的.pt文件
–source：测试图片的路径

这样就完成了