YOLOv5-6.1剪枝实战
前言
为降低网络模型的计算量、减小计算和训练时间、以及降低内存占用,特别是在性能较差的网络设备上的应用,出现了一系列模型简化方法。网络剪枝、量化、低秩分解、知识蒸馏等是其中的代表。本文将重点介绍网络剪枝这一简化方法。
这里引用51CTO白老师一个很形象的例子,人从出生到六岁大脑神经元是在增加的,但是从幼年到青年神经元的数量在减少,一些权重很小甚至不必要的连接在被取消。通过这个例子可以对剪枝有一个初步的印象。
YOLO版本的确定
无论是正常训练还是剪枝轻量化还是模型转换的部署工作,在大家选择任何一个教程的时候都要注意yolo的版本,即使是大版本v5统一的前提下也要看清楚是5.0、6.1还是6.2,这里剪枝的教程选择的是YOLOv5-6.2,后续将在部署的章节详细记录YOLO如何甄别版本号以及如何去github上下载正确的版本,这里不再赘述。
什么是剪枝?
Learning Efficient Convolutional Networks Through Network Slimming,这里给大家推荐一个很经典的论文,如果能力稍微差一些的同学可以选择51CTO这种的网课平台,或者B站等免费公开课资源,在正式实践之前至少要搞明白剪枝在做什么。
简单描述为:在批归一化(BN)层中,存在两个可训练参数,即γ(gamma)和β(beta)。当γ趋于0时,输入经过BN后相当于乘以了0,导致该通道上的卷积输出为0,这是不希望看到的情况。因此,可以认为去除这种冗余通道对模型性能不会产生影响。通常情况下,由于初始化,γ的值约为1。为了将γ趋近于0,可以通过引入L1正则化来约束,以实现系数的稀疏化。文献中将添加γ的L1正则化称为稀疏训练。
剪枝实战
在剪枝开始之前依旧可以进行自制数据集的标注或者开源数据集的下载划分,然后编写配置文件等工作,也就是说剪枝之前需要进行一次正常完整的YOLO训练。这里需要强调一点就是大多数剪枝都是在进行在线式的剪枝,也就是,而边训练边剪枝,这种方式有一种更加广泛的称呼,即稀疏化训练。而离线式剪枝我这里推荐一个写的特别好的博主的文章离线剪枝。
稀疏化训练
python train_sparsity.py --st --sr 0.005 --data data/yours.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 100 --workers 4 --name yolov5s-sparse
sr的选择需要根据数据集调整,可以先设置0.001-0.005进行尝试,如果你的数据集比较庞大可以尝试适当增大你的数值。
这里有个很重要的一点就是需要对训练过程进行可视化,对map的变化和bn训练中缩放因子分布的变化反过来对sr的数值进行调整。
如果你是本地显卡训练,则在对应的exp中使用下面这条命令
tensorboard --logdir=./runs/train/yolov5s-sparse
然后根据提示浏览器打开http://localhost:6006/
如果你是在线服务器例如AutoDL,则使用下面这条命令
首先结束默认启动的TensorBoard进程,执行命令:ps -ef | grep tensorboard | awk '{print $2}' | xargs kill -9
在终端中执行以下命令启动TensorBoard $ tensorboard --port 6007 --logdir /path/to/your/tf-logs/direction
训练完成后在对缩放因子进行可视化
调整的原则就是map不要下降太多,然后还可以使横轴的峰值应随着训练次数不断逼近0轴。如果峰值逼近x轴的速度缓慢,形象的说法就是越往下(越靠近x轴)山峰不向着0轴靠近,那么sr的值就应该适量的调大。
网络模型剪枝
这里要使用稀疏化训练后的last.pt
python prune.py --percent 0.5 --weights runs/train/yolov5s-sparse/weights/last.pt --data data/yours.yaml --cfg models/yolov5s.yaml
--percent 0.5是剪裁的比率
运行之后会显示剪枝前后的通道对比
剪枝以后,在我这个示例当中精度下降很多。但是大小从27.3下降到10.9MB,这里二者依旧是32位浮点。
剪枝后的模型微调
python finetune_pruned.py --weights pruned_model.pt --data data/yours.yaml -- batch-size 16 --workers 4 --epochs 100 --name yolov5s-finetuned
稀疏化训练=>剪枝=>微调 可以反复迭代这个过程多次进行剪枝。
更多推荐
所有评论(0)