从0开始做yolov5模型剪枝

从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网络