从0开始做yolov5模型剪枝 ·
1 前言
【整个流程中，在正常train，sparityTrain，prune，finetune遇到10多个的问题，包括AttributeError、ModuleNotFoundError、RuntimeError、SyntaxError、TypeError等问题的解决方法，详见内容】
为了将现有模型移植到ARM平台，同时保证模型准确率的基础上，减少模型的算力消耗和推理时间。
之前有做实验对比了YOLOv5、 YOLOv7、 YOLOv8，结合不同版本模型推理时间和准确率，并查了很多资料，包括大部分人博客描述，结合大部分人经验，我们觉的yolov5的泛化能力较好。故在考虑训练自己的模型且在X86和ARM平台上部署，我们针对yolov5做模型的训练和剪枝，以便轻量化小模型的部署。
当然，我们还需要对最终的模型执行INT8量化的操作，以便降低目标检测的推理时间。
2 GitHub取源码
下载如下路径的源码：
https://github.com/midasklr/yolov5prune/tree/v6.0
本文为 上面GitHub上取6.0的版本做剪枝
3 原理
【根据一些博客/文章对yolov5剪枝的介绍，简单总结一下yolov5模型剪枝的原理】
3.1 原理
原理论文：Learning Efficient Convolutional Networks through Network Slimming
ref: Pruning Filters for Efficient ConvNets( https://arxiv.org/abs/1608.08710 )
ref: https://blog.csdn.net/qq_42835363/article/details/129125376?spm=1001.2014.3001.5501
ref: https://blog.csdn.net/IEEE_FELLOW/article/details/117236025
ref: Yolov5_5.0上做模型剪枝
输入经过BN（Batch Normalization）层获得归一化后的分布。BN层存在两个可训练参数γ(gamma)、β(beta)。
当gamma和betaγ趋于0时，输入相当于乘以了0，此时该channel上的卷积将输出0，这是无意义的。因此，可以认为剔除这样冗余的channel对模型性能没有影响。
普通网络训练时，由于初始化，gamma一般分布在1附近。为了使gamma趋于0，可以通过添加L1正则来约束，使得系数稀疏化。论文中把添加gammaL1正则的训练称为稀疏训练。
稀疏训练后，裁剪掉稀疏很小的层，对应激活也很小，所以对后面的影响非常小，反复迭代这个过程，可以获得小型的模型，步骤如图1。
图1
3.2 network slimming过程
① 先初始化网络，对BN层的参数添加L1正则并对网络训练。
② 统计网络中的γ(gamma)，设置剪枝率对网络进行裁剪。
③ 将裁减完的网络finetune，完成剪枝工作。
4 具体实施步骤
4.1 安装虚拟环境
解压下载的源码，进入yolov5prune_6.0目录下，依次执行下面的操作
# 1 创建虚拟环境
conda create -n yolov5prune
# 2 激活虚拟环境
conda activate yolov5prune
# 3 安装虚拟环境（根据yolov5prune_6.0根目录下的requirements.txt安装）
pip install -r requirements.txt
4.2 配置参数
4.2.1 数据集参数
自己的数据集结构如下
--datasTrain
------images
----------train # 存放训练数据集的图片（.jpg）
----------val
----------test
------labels
----------train # 存放训练图片对应的标签文件（.txt）
----------val
----------tes
在/yolov5prune_6.0/data/目录下，仿照coco128.yaml中的结构创建my_yolov5.yaml文件。其中内容如下
# Train/val/test sets as
# 1) dir: path/to/imgs,
# 2) file: path/to/imgs.txt, or
# 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: /home/user/hlj/MyTrain/datasTrain3_More/ # dataset root dir
train: images/train/ # train images (relative to 'path') 128 images
val: images/val/ # val images (relative to 'path') 128 images
test: images/test/ # test images (optional)
nc: 11 # number of classes
names: ['pedes', 'car', 'bus', 'truck', 'bike', 'elec', 'tricycle', 'coni', 'warm', 'tralight', 'specVehi']
4.2.2 模型结构参数
修改yolov5prune_6.0/models/yolov5s.yaml中的目标检测类型，使其适配自己数据集的目标检测类型数量。如下
nc: 11
4.2.3 train.py中的参数
设置train.py中的参数，主要包括如下：
'--weights', default='./yolov5s.pt' # 由于我要从头训练，所以注释了此参数
'--cfg', default='./models/yolov5s.yaml'
'--data', default='./data/my_yolov5.yaml'
'--epochs', default=300 # 由于从头训练，所以epochs值设的比较大
'--batch-size', default=-1
'--imgsz', default=640 # 考虑部署
4.3 正常训练
4.3.1 准备
由于我是SSH链接，所以先创建/打开tmux会话
tmux new -s prunesession
若【先按下ctrl+b，然后再单独按d】退出会话的话，下次再进入会话，需要使用命令
tmux a -t prunesession
进入会话，先进入项目目录下，并激活虚拟环境(若已激活，可忽略)
cd ../yolov5prune_6.0/
source activate yolov5prune
训练结束之后要删除会话
tmux kill-session -t prunesession
4.3.2 训练及问题解决
执行如下命令，进行训练
python3 train.py
【问题1】
运行train.py文件后，报了如下错误
ModuleNotFoundError: No module named 'utils.loggers.wandb'
提示缺包，根据别人的攻略，下载U神对应yolov5_6.0版本的代码，然后把yolov5_6.0\utils\loggers\目录下的整个wandb文件夹拷贝到yolov5prune_6.0\utils\loggers目录下。
【问题2】
重新输入 python3 train.py 后，报如下问题，由此可见，设置train.py中的参数的时候，'–weights’参数不可以被注释掉。
AttributeError: 'Namespace' object has no attribute 'weights'
故将’weights’参数设置如下，表示不使用预训练权重，模型将从头开始训练。
'--weights', default=''
【问题3】
不知道为什么，ubuntu上又报了一个numpy的问题如下，本地运行是没有这个问题的
raise AttributeError(__former_attrs__[attr])
AttributeError: module 'numpy' has no attribute 'int'.
原来是因为 新版本的numpy里面没有np.int，可以修改源码解决。
修改yolov5prune_6.0/utils/ 目录下datasets.py中所有的…astype(np.int) 为 …astype(int)，如下所示：
441 bi = np.floor(np.arange(n) / batch_size).astype(int) # batch index
483 self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(int) * stride
854 b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(int)
修改yolov5prune_6.0/utils/ 目录下general.py中所有的…astype(np.int) 为 …astype(int)，如下所示：
510 classes = labels[:, 0].astype(int) # labels = [class xywh]
525 class_counts = np.array([np.bincount(x[:, 0].astype(int), minlength=nc) for x in labels])
【问题4】
File "/home/user/hlj/MyTrain/yolov5prune_6.0/utils/loss.py", line 217, in build_targets
indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))
RuntimeError: result type Float can't be cast to the desired output type long int
参考 ref: https://blog.csdn.net/Thebest_jack/article/details/125649451 执行如下操作：
修改yolov5prune_6.0/utils/ 目录下loss.py源码，
#（1） 183行左右
for i in range(self.nl):
anchors, shape = self.anchors[i], p[i].shape # anchors = self.anchors[i]
gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain
#（2）218行后
# indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))
上一行代码改为如下
indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1))) # image, anchor, grid indices
【问题5】
epoch的时候，报如下问题
File "..../yolov5prune_6.0/utils/plots.py", line 116, in text
w, h = self.font.getsize(text) # text width, height
AttributeError: 'FreeTypeFont' object has no attribute 'getsize'
这是因为安装了新版本的 Pillow，pip install tf-models-official删除了该getsize功能
，降级到 Pillow 9.5 解决了该问题。可以尝试以下方法进行解决，参见【问题8】：
pip install Pillow==9.5
【问题6】
epoch 0 结束，val结束后，报如下问题
File ".....\yolov5prune_6.0\utils\callbacks.py", line 77, in run
logger['callback'](*args, **kwargs)
TypeError: on_fit_epoch_end() missing 1 required positional argument: 'fi'
找一下官方源码，把yolov5_6.0/utils/下的整个loggers文件复制进去，应该就行了，可能是版本不一致导致的。
【问题7】
yolov5prune_6.0/utils/general.py line471
return re.sub(pattern="[|@#!?·$€%&()=??^*;:,¨′><+]", repl="_", string=s)
SyntaxError：(unicoda error)'utf-8' code can't decode byte 0xal in position 6: invalid start byte。
应该是’utf-8’ code不支持的问题，添加了下面的编码格式，但最后也没解决。看了一下对应函数的功能，只是为了清理字符串（用下划线替代特殊字符），所以就直接改了那一行的代码，对整个程序是没有影响的。
# -*- coding: utf-8 -*-
【总结】
虽然各种问题不断，而且【问题5】我也没管它 但是，最后总算是python3 train.py正常跑起来了。
4.4 稀疏化训练
4.4.1 参数设置
对train_sparity.py的参数进行设置
'--st', action='store_true',default=True,
'--sr', type=float, default=0.0001,
'--weights', type=str, default=ROOT / '',
'--cfg', type=str, default='./models/yolov5s.yaml',
'--data', type=str, default='./data/my_yolov5.yaml',
'--epochs', type=int, default=300
'--batch-size', type=int, default=-1, # 注意【问题8】的发生
'--imgsz', '--img', '--img-size', type=int, default=640,
'--adam', action='store_true', default=True,
4.4.2 稀疏化训练与问题
执行如下命令，进行稀疏化训练
python train_sparity.py
【问题8】
loggers.on_params_update({"batch_size": batch_size})
AttributeError: 'Loggers' object has no attribute 'on_params_update'
貌似是autobatch的原因，所以把参数’–batch-size’, type=int, default=-1, 先改为固定值 default=2，此后，可以正常epoch0。但是仍然存在【问题5】的问题，虽然不影响训练，但是觉的还是应该把它解决，毕竟是AttributeError的问题。解决办法如下：
# pillow版本太新的原因，新版的getsize属性被删除掉了。
pip3 uninstall pillow
pip3 nstall pillow==9.5
【问题9】
Epoch0的val结束之后，报了如下问题
File "/home/user/hlj/MyTrain/yolov5prune_6.0/utils/callbacks.py", line 77, in run
logger['callback'](*args, **kwargs)
TypeError: Loggers.on_fit_epoch_end() takes 5 positional arguments but 6 were given
这个问题是由于自己为了解决【问题6】，把项目中的utils/loggers/init .py文件换成了官方的文件了，发现init.py文件中def on_fit_epoch_end(self, vals, bn_weights, epoch, best_fitness, fi)少bn_weights，在这个项目中把prune项目下这个文件重新拷贝一下好了。
4.5 剪枝
4.5.1 参数设置
设置裁剪比例参数，可以从小到大试。注意cfg的模型文件需要和weights对应上，否则会出现运行prune 过程中出现键值不对应的问题，裁剪完成会保存对应的模型pruned_model.pt。
在prune.py文件中，修改如下参数
'--data', type=str, default=ROOT / 'data/my_yolov5.yaml',
'--weights', nargs='+', type=str, default=ROOT / 'runs/train/spaweight/last.pt'
'--cfg', type=str, default='./models/yolov5s.yaml',
'--percent', type=float, default=0.1,
'--batch-size', type=int, default=16,
'--imgsz', '--img', '--img-size', type=int, default=640,
运行
python prune.py
【问题10】
SyntaxError: Non-UTF-8 code starting with '\xe5' in file /home/user/hlj/MyTrain/yolov5prune_6.0/prune.py on line 400, but no encoding declared; see https://peps.python.org/pep-0263/ for details
解决方法：找到对应的行，发现是注释的内容code的格式的问题，把它删掉或者把中文改成英文即可。
【问题11】
return func(*args, **kwargs)
TypeError: run() got an unexpected keyword argument 'cfg'
解决方式在prune.py的源码run()函数中，增加参数如下
cfg = './model/yolov5s.yaml'
4.5.2 剪枝
对稀疏化训练后的模型best.pt进行剪枝。
若参数已设置好了，直接执行python prune.py
python prune.py
否则，传入的weight为稀疏化训练得到的权重。
python prune.py --weights runs/train/exp_sparity/weights/best.pt --percent 0.5 --cfg models/yolov5s.yaml
裁剪完成会在根目录下保存对应的模型pruned_model.pt。
4.6 finetune剪枝的网络
4.6.1 参数设置
更改finetune_pruned.py的相关参数如下
'--weights', type=str, default=ROOT / 'pruned_model.pt',
'--cfg', type=str, default='./models/yolov5s.yaml',
'--data', type=str, default=ROOT / 'data/my_yolov5.yaml',
'--epochs', type=int, default=100
'--batch-size', type=int, default=16,
'--imgsz', '--img', '--img-size', type=int, default=640,
'--adam', action='store_true', default=True,
'--workers', type=int, default=8,
'--project', default=ROOT / 'runs/finetune',
4.6.2 finetune
若没有修改finetune_pruned.py中的参数，如下执行。
python finetune_pruned.py --weights pruned_model.pt --adam --epochs 100
由于直接修改了finetune_pruned.py中的参数，直接执行
python finetune_pruned.py
执行时，报了【问题9】，按照相关解决方法可fine_tune正常开始。
4.7 循环稀疏训练->剪枝->finetune网络
此次仅做了一轮的稀疏化->剪枝->微调。其中正常训练和稀疏化训练均没有使用预训练模型。对比包含针对模型不同尺度（10%、20%、30%）的裁剪后finetune的结果。
表2 目标检测模型剪枝对比
| Train | Sparity | Before Prune | fintune | fintune | fintune | |
|---|---|---|---|---|---|---|
| Prune0.1 | Prune0.2 | Prune0.3 | ||||
| modelSize(M) | 13.8M | 13.8M | 13.8 | 12.3M | 10.5M | 8.68M |
| P | 0.915 | 0.923 | 0.919 | 0.916 | 0.899 | 0.892 |
| R | 0.832 | 0.808 | 0.816 | 0.829 | 0.84 | 0.842 |
| [email protected] | 0.885 | 0.879 | 0.881 | 0.887 | 0.887 | 0.888 |
| [email protected]:.95 | 0.642 | 0.628 | 0.633 | 0.639 | 0.64 | 0.642 |
（1）可以看出，保证基本mAP的情况下，对yolov5s的模型剪枝30%是可行的。
（2）考虑模型的针对性，通常我们很难训练一个准确率很高，泛化能力又很强的模型，所以往往我们需要针对项目数据及模型进行剪枝和微调。
（3）对微调后的模型测试检测实时视频流的效果，鉴于模型训练尺度是640，存在远处目标召回率较低，路口流量密集时，存在漏检现象。