UltraFastYoloV5发布!基于剪枝技术的模型压缩,3ms推理的YoloV5mAP60+!

UltraFastYoloV5发布!基于剪枝技术的模型压缩,3ms推理的YoloV5mAP60+!

本文由林大佬原创，转载请注明出处，来自腾讯、阿里等一线AI算法工程师组成的QQ交流群欢迎你的加入: 1037662480

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

• 开源的模型根本没法用, 只能做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的模型体积. 这对于移动端来说是非常友好的, 如果你想要更快的速度, 那么减小输入尺寸, 还可以获得加一倍的速度!

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

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

裁剪之前的准备

假如你关注了我们的专栏, 我们之前有写过一篇自己构建Yolov5模型的文章. 官方的Yolov5是没有办法进行裁剪的, 因为它的构建方式有点复杂, 很难对模型进行分析.

建议大家可以看我们之前的文章:

https://zhuanlan.zhihu.com/p/242456389

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

代码

Talk is cheap, 幸运的是我们所有的操作都在Yolov5的fork版本之上完成, 你可以很方便看到我们裁剪之后的模型, 甚至可以以和yolov5一模一样的训练方式,复现我们的结果. 大家可以在我们的平台找到对应的代码, 下载下来进行尝试:

http://manaai.cn/aicodes_detail3.html?id=66

更多

如果你想学习人工智能，对前沿的AI技术比较感兴趣，可以加入我们的知识星球，获取第一时间资讯，前沿学术动态，业界新闻等等！你的支持将会鼓励我们更频繁的创作，我们也会帮助你开启更深入的深度学习之旅！

往期文章

https://zhuanlan.zhihu.com/p/165009477

https://zhuanlan.zhihu.com/p/149398749

https://zhuanlan.zhihu.com/p/147622974

https://zhuanlan.zhihu.com/p/144727162