YoloV5 最强剪枝技术! 模型压缩,3ms推理的YoloV5mAP60+!

喜欢ヅ旅行 2022-10-16 05:00 528阅读 0赞

先看效果图：

![1374e5944ef2157fd31fd14ddc46f3db.gif][]![wAAACH5BAEKAAAALAAAAAABAAEAAAICRAEAOw==][wAAACH5BAEKAAAALAAAAAABAAEAAAICRAEAOw]

![6524c608e4f65e9d7829aaa959f3d1fb.gif][]![wAAACH5BAEKAAAALAAAAAABAAEAAAICRAEAOw==][wAAACH5BAEKAAAALAAAAAABAAEAAAICRAEAOw]

由于我们修改了model, 没有加pretrain的模型, 但是即便如此, 3个epoch就到了 68.9 的mAP!

前段时间受疫情的影响, 公司希望开发一个口罩佩戴检测的应用, 讲道理这很简单, 我记得在年初的时候有一个团队就发布了这么一个模型, 但我们尝试下来发现事情并没有那么简单, 原因是:

开源的模型根本没法用, 只能做demo,实际场景下漏检非常严重;  
速度并不快. 很多人说30FPS不就已经够了吗? 公司给我们的就一个JetsonNano的开发板, 再快的模型跑到这上面都要凉凉.  
因此, 一个很严峻的问题摆在我们面前, 如果基于现有的最牛逼的检测器开发一个足够快, 模型体积足够小, 同时精度又还OK的模型呢? 我们得出的解决方案只有一个:对Yolov5进行模型剪枝和量化！

说到这里, 实际上我们做的并没有那么透彻, 因为这不是一个完完全全的剪枝套路, 我们做的工作主要有三个:

原始的YoloV5的backbone过于重, 先用一个小的backbone做一个baseline;  
对这个backbone的HEAD部分进行裁剪;  
对Detect Layer进行裁剪.  
对整个模型进行剪枝计算, 得到剪枝表, 估算我们要裁剪的部分, 得到裁剪之后的模型.  
很多同学应该都有试过Yolov5的模型, 它的检测能力还是非常强的, 鄙人认为主要体现在它的结构除了包含Yolov4的所有优点以外, 相较于Yolov4训练过程更加的科学, 对于正样本的挖掘和有效训练可以说发挥到了极致. 另外不同数据集对于Anchor的作用是非常敏感的, 知乎也有大佬分析过这个因素对于最终训练结果的影响, 本人也做过实验, 一个无比正确的anchor决定了你的检测效果的80%,. 这也是为什么很多SSD系列的算法在不同数据集上有些人就是得不到好的结果 (包括EfficientDet). 因为缺乏一个正确的anchor的指导.

但本篇文章想阐述的是, 即便是Yolov5, 它也存在着极大的参数冗余. 这也是我们进行剪枝的一个出发点. 实验结果上经过我们的裁剪, Yolov5s的模型大小可以降低 29倍. 表格对比如下:

YoloV5s    Yolov5s-Pruned  
model size    29M    1.3M  
speed (TensorRT GPU)    12ms    4ms  
CPU (ONNXRuntime)    400ms    30ms  
这个数据所有的输入尺寸都为 1280x800. 最终我们的效果可以做到 4ms 1.3M的模型体积. 这对于移动端来说是非常友好的, 如果你想要更快的速度, 那么减小输入尺寸, 还可以获得加一倍的速度!

我们测试数据集就是对口罩进行检测数据集, 实际效果可以看到:

![620d31d0f9cd43cc8e50c223a1896dde.png][]

对于小目标的检测效果还是非常不错的! 尤其是这是图片都是来自于没有被训练的百度图片, 说明模型的鲁棒性也还是OK的! 事实说明, 我们的裁剪策略是非常有效的, 对于一些面向移动端的场合, 这几乎是必经之路.

由于裁剪部分的操作还是有点复杂, 而且有点杂乱, 我们会在后续的文章中陆续教授,也可能以视频录制的形式开发出来,请大家关注我们的专栏, 或者我们的公众号来接受第一时间的推送.

代码：

代码可以联系我：

![2021042216121945.jpg][]

[1374e5944ef2157fd31fd14ddc46f3db.gif]: https://img-blog.csdnimg.cn/img_convert/1374e5944ef2157fd31fd14ddc46f3db.gif
[wAAACH5BAEKAAAALAAAAAABAAEAAAICRAEAOw]: 
[6524c608e4f65e9d7829aaa959f3d1fb.gif]: https://img-blog.csdnimg.cn/img_convert/6524c608e4f65e9d7829aaa959f3d1fb.gif
[620d31d0f9cd43cc8e50c223a1896dde.png]: https://img-blog.csdnimg.cn/img_convert/620d31d0f9cd43cc8e50c223a1896dde.png
[2021042216121945.jpg]: https://img-blog.csdnimg.cn/2021042216121945.jpg