华为P30上用NCNN移植跑benchmark,每帧仅需10.23毫秒,比yolov4-tiny快3倍,参数量小6倍,COCO mAP(0.5:0.95)能够达到20.6 。而且模型权重文件只有1.8mb,对比动辄数十兆的模型,可以说是相当友好了~
Android Demo
项目地址(提供训练代码到安卓部署一站式解决方案):
RangiLyu/nanodet: ⚡Super fast and lightweight anchor-free object detection model. Only 1.8mb and run 97FPS on cellphone (github.com)github.com
深度学习目标检测已经发展了许多年,从Two-stage到One-stage,从Anchor-base到Anchor-free,再到今年的用Transformer做目标检测,各种方法百花齐放,但是在移动端目标检测算法上,yolo系列和SSD等Anchor-base的模型一直占据着主导地位。这个项目的主要目的是希望能够开源一个移动端实时的Anchor-free检测模型,能够提供不亚于yolo系列的性能,而且同样方便训练和移植。
其实从去年大量anchor-free的论文发表之后,我就一直想把anchor free的模型移植到移动端或者是嵌入式设备上。当时尝试过把FCOS轻量化的实验,奈何效果并不如mobilenet+yolov3,也就暂时搁置了。分析下来,主要还是因为FCOS的centerness分支在轻量级的模型上很难收敛,之后发表的一些在FCOS上面进行改进的论文也都没有解决这一问题。
直到今年年中的时候刷arxiv突然刷到了
@李翔
老师的论文Generalized Focal Loss: Learning Qualified and Distributed Bounding Boxes for Dense Object Detection,论文中提出的GFocalLoss完美去掉了FCOS系列的Centerness分支,而且在coco数据集上涨点显著,这么好的文章怎么能不点赞呢?GFL的出现不仅去掉了难以训练的Centerness,而且还省去了这一分支上的大量卷积,减少了检测头的计算开销,非常适合移动端的轻量化部署。
贴一下李翔老师关于GFocal Loss的解读:
李翔:大白话 Generalized Focal Losszhuanlan.zhihu.com
在找到了合适的损失函数之后,接下来的关键就是如何使其在轻量化模型上发挥作用了。首先需要对移动端进行优化的就是检测头:FCOS系列使用了共享权重的检测头,即对FPN出来的多尺度Feature Map使用同一组卷积预测检测框,然后每一层使用一个可学习的Scale值作为系数,对预测出来的框进行缩放。
FCOS模型架构
这么做的好处是能够将检测头的参数量降低为不共享权重状态下的1/5。这对于光是检测头就拥有数百通道的卷积的大模型来说非常有用,但是对于轻量化模型来说,共享权重检测头并没有很大的意义。由于移动端模型推理由CPU进行计算,共享权重并不会对推理过程进行加速,而且在检测头非常轻量的情况下,共享权重使得其检测能力进一步下降,因此还是选择每一层特征使用一组卷积比较合适。
同时,FCOS系列在检测头上使用了Group Normalization作为归一化的方式,GN对比BN有很多好处,但是却有一个缺点:BN在推理时能够将其归一化的参数直接融合进卷积中,可以省去这一步计算,而GN则不行。为了能够节省下归一化操作的时间,我选择将GN替换为BN。
FCOS的检测头使用了4个256通道的卷积作为一个分支,也就是说在边框回归和分类两个分支上一共有8个c=256的卷积,计算量非常大。为了将其轻量化,我首先选择使用深度可分离卷积替换普通卷积,并且将卷积堆叠的数量从4个减少为2组。在通道数上,将256维压缩至96维,之所以选择96,是因为需要将通道数保持为8或16的倍数,这样能够享受到大部分推理框架的并行加速。最后,借鉴了yolo系列的做法,将边框回归和分类使用同一组卷积进行计算,然后split成两份。下图就是最终得到的轻量化检测头的结构,非常的小巧:
NanoDet检测头
目前针对FPN的改进有许多,如EfficientDet使用了BiFPN,yolo v4和yolo v5使用了PAN,除此之外还有BalancedFPN等等。BiFPN虽然性能强大,但是堆叠的特征融合操作势必会带来运行速度的降低,而PAN只有自上而下和自下而上两条通路,非常的简洁,是轻量级模型特征融合的好选择。
原版的PAN和yolo中的PAN,都使用了stride=2的卷积进行大尺度Feature Map到小尺度的缩放。我为了轻量化的原则,选择完全去掉PAN中的所有卷积,只保留从骨干网络特征提取后的1x1卷积来进行特征通道维度的对齐,上采样和下采样均使用插值来完成。与yolo使用的concatenate操作不同,我选择将多尺度的Feature Map直接相加,使得整个特征融合模块的计算量变得非常非常小。
最终的极小版PAN的结构也是非常简单:
超轻量的PAN
原本有考虑过自己魔改一个轻量级的backbone出来,但是最后评估了一下感觉工作量太大了(在家训练模型电费太贵),就打算选择使用现有的一些轻量级骨干网络。最开始的选择有MobileNet系列,GhostNet,ShuffleNet,还有最近的EfficientNet。在评估了参数量、计算量以及权重大小之后,还是选择了ShuffleNetV2作为骨干网络,因为它是这些模型里面在相似精度下体积最小的,而且对移动端CPU推理也比较友好。
最终我选择使用ShuffleNetV2 1.0x作为backbone,去掉了最后一层卷积,并且抽取8、16、32倍下采样的特征输入进PAN做多尺度的特征融合。整个骨干模型使用了Torchvision提供的代码,能够直接加载Torchvision上提供的imagenet预训练权重,对加快模型收敛有很大的帮助。顺便一提,最近有些论文指出使用分类的预训练权重初始化模型对检测任务的效果不如随机初始化的好,不过这要付出训练更多步数的代价,我还没有测试过,欢迎大家尝试~
NanoDet整体模型结构
在经过对one-stage检测模型三大模块(Head、Neck、Backbone)都进行轻量化之后,得到了目前开源的NanoDet-m模型,在320x320输入分辨率的情况下,整个模型的Flops只有0.72B,而yolov4-tiny则有6.96B,小了将近十倍!模型的参数量也只有0.95M,权重文件在使用ncnn optimize进行16位存储之后,只有1.8mb,非常适合在移动端部署,能够有效减少APP体积,同时也对更低端的嵌入式设备更加友好。
尽管模型非常的轻量,但是性能却依旧强劲。对于小模型,往往选择使用AP50这种比较宽容的评价指标进行对比,这里我选择用更严格一点的COCO mAP(0.5:0.95)作为评估指标,同时兼顾检测和定位的精度。在COCO val 5000张图片上测试,并没有使用Testing-Time-Augmentation的情况下,320分辨率输入能够达到20.6的mAP,比tiny-yolov3高4分,只比yolov4-tiny低1个百分点,而将输入分辨率与yolo保持一致,都使用416输入的情况下,得分持平。
最后用ncnn部署到手机上之后跑了一下benchmark,模型前向计算时间只要10毫秒左右,对比yolov3和v4 tiny,均在30毫秒的量级。在安卓摄像头demo app上,算上图片预处理,检测框后处理以及绘制检测框的时间,也能轻松跑到40+FPS~。(关于模型从pytorch到安卓的部署,以及如何用C++实现GFL检测框的后处理,我会再专门发一遍文章详细介绍)
ncnn 安卓benchmark对比
NanoDet总体而言没有特别多的创新点,是一个纯工程化的项目,主要的工作就是将目前学术界的一些优秀论文,落地到移动端的轻量级模型上。最后通过这些论文的组合,得到了一个兼顾精度、速度和体积的检测模型。
为了让大家能够快速使用NanoDet,方便训练和部署,我已经把Pytorch的训练代码、基于NCNN的Linux和windowsC++部署代码以及Android摄像头Demo全部开源,同时在Readme中也包含了非常详细的教程,欢迎大家使用,欢迎提issue~
顺便提一句,NanoDet训练并没有用到很多数据增强的技巧,模型结构也非常简单,也就是说mAP应该还有很大的提升空间,如果有小伙伴愿意魔改一下再涨几个点那就再好不过了。