从0开始做yolov5模型剪枝

文章目录

  • 从0开始做yolov5模型剪枝 **·**
    • 1 前言
    • 2 GitHub取源码
    • 3 原理
      • 3.1 原理
      • 3.2 network slimming过程
    • 4 具体实施步骤
      • 4.1 安装虚拟环境
      • 4.2 配置参数
        • 4.2.1 数据集参数
        • 4.2.2 模型结构参数
        • 4.2.3 train.py中的参数
      • 4.3 正常训练
        • 4.3.1 准备
        • 4.3.2 训练及问题解决
      • 4.4 稀疏化训练
        • 4.4.1 参数设置
        • 4.4.2 稀疏化训练与问题
      • 4.5 剪枝
        • 4.5.1 参数设置
        • 4.5.2 剪枝
      • 4.6 finetune剪枝的网络
        • 4.6.1 参数设置
        • 4.6.2 finetune
      • 4.7 循环稀疏训练->剪枝->finetune网络

从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。

从0开始做yolov5模型剪枝_第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

从0开始做yolov5模型剪枝_第2张图片

【问题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网络

你可能感兴趣的:(YOLO,剪枝,算法)