工程Gitee地址:
https://gitee.com/zhong-liangtang/ncnn-android-yolox-nano
一、YOLOX简介
YOLOX是一个在2021年被旷视科技公司提出的高性能且无锚框(Anchor-free)的检测器,在YOLO系列的基础上吸收近年来目标检测学术界的最新成果,如解耦头(Decoupled Head)、数据增强、无锚框、标签分配策略SimOTA(Simplified Optimal Transport Assignment)等等。同时YOLOX继承了YOLO系列容易部署的特点,提供了支持ONNX、TensorRT、ncnn和Openvino的部署版本。
二、YOLOX网络结构
YOLOX的网络有多种版本,本文介绍YOLOX-Nano的网络结构,它由四个部分构成,分别为输入端、用于特征提取的主干网络(Backbone)、用于特征融合的颈部网络(Neck)以及预测(Prediction)。YOLOX-Nano模型网络结构图如下所示。
输入端采取了Mosaic和Mixup两种数据增强方法,在输入图片时将图片的高和宽统一进行等比例转换。
主干网络使用CSPDarknet网络进行特征提取,CSPDarknet网络利用52层卷积神经网络进行特征提取,1层利用1×1卷积作为全连接层,共计53层卷积神经网络。在特征提取时,会把在特征提取过程中三部分不同尺度的特征信息直接和下一部分的颈部网络相连接。
颈部网络使用FPN结构进行特征融合,在该部分,已经在主干网络获得的特征会结合不同尺度的特征信息来进一步进行特征提取,使模型的性能以及小目标的检测能力得到提升。
在预测部分,在该部分使用解耦头、无锚框、标签分配策略SimOTA等先进技术,最后将所有的特征信息统一转换成二维矩阵[16]。
三、 YOLOX主要改进部分
1、解耦头
为解决分类和回归冲突问题[8],YOLOX提出了解耦头这一解决方案。YOLO系列所采用的耦合头(Coupled Head)把分类和回归任务都放在一个1×1卷积里实现,它首先通过一个1×1卷积把前面的特征图的通道数统一进行调整,再经过2个3×3卷积层来提取特征,最后利用一个1×1卷积层,分别到分类和回归检测头,同时在回归分支还增加了IoU分支,解耦头工作流程图如图所示。在使用了解耦头之后,训练时收敛速度和准确率都有所提高,并且解耦头对端到端版本的YOLO会提高一定的准确率,这些都体现了解耦头的训练和预测价值。
2、数据增强
YOLOX采用了Mosaic和MixUp数据增强方法。
MixUp数据增强方法是通过线性插值的方式混合两个样本和标签,扩大了训练数据,同时提升了模型的泛化能力。
Mosaic数据增强方法是在训练期间将四张图进行裁剪、拼接成一张新图,大幅度扩充了图片数据,避免因训练集背景相似而降低模型的泛化能力。
然而,轻量化模型YOLOX-Nano如果同时使用这两种数据增强的方法,在COCO数据集上的AP反而会降低,但在不使用MixUp数据增强方法的同时削弱Mosaic数据增强方法的使用,AP反而有所上升。
3、无锚框
YOLOX是一个基于无锚框的检测器,而之前的YOLO系列都是基于锚框(Anchor-based)的检测器,但基于锚框的检测器存在较多的问题。首先,锚框大小的选择是采用聚类分析的方法,在训练之前可以得到比较好的锚框,但是这些聚类的锚框没有一定的通用性,在预测新的图片时效果不佳。其次,锚框机制增加了检测头的复杂度,并且对每张图片增加了预测的数量,因此也增加了预测的时间[20]。由于不使用锚框进行预测与使用锚框相比在COCO数据集上的AP基本相同,因此使用无锚框策略对于YOLOX检测器的预测速度和准确率都有所提升。
4、标签分配策略SimOTA
SimOTA是OTA的简化版本,OTA是解决候选框分配问题,一般在分配正负样本时是根据ground-truth的IoU与Anchor进行分配,但是正负样本的分配需要优先考虑全局[21]。所以我们通常把候选框分配问题当成一个线性规划中的优化传输问题(Optimal Transport,OT)来处理,其原理是建立一个代价矩阵,如果有M个ground-truth以及与其对应的N个候选框,即M×N为代价矩阵, 矩阵中的每个元素就是该ground-truth与候选框的Loss值,Loss值越大则说明选取这对ground-truth和候选框的代价越大,优化传输的目的是选取ground-truth与候选框相匹配,使总体代价最小。
5、YOLOX-Nano模型
轻量化模型YOLOX-Nano是所有YOLOX模型中参数量最小的模型,其参数量(Parameters)为0.9M,浮点运算数(FLOPs)为1.08G,而它比参数量相当的NanoDet模型在COCO数据集上的AP高1.8%,比YOLOv4-Tiny模型的AP高10%[8]。更为重要的是,轻量化模型YOLOX-Nano可以被部署到移动端或其他嵌入式设备上,说明它有一定的工业价值,在应对实时检测的同时也能与设备的其他功能进行配合,共同完成工作。
四、试验环境与模型训练
模型训练的GPU为NVIDIA GeForce RTX 3070 Laptop GPU 8GB,在基于安装Anaconda 4.10.3和CUDA 11.1的Windows 10操作系统下进行训练,采用cuDNN 8.0深度学习库,处理器为AMD Ryzen 7 5800H with Radeon Graphics,计算机内存为16G,使用开源深度学习框架Pytorch作为开发环境。
模型训练方式及模型评价指标
本文利用采集的数据集训练轻量化模型YOLOX-Nano,使用YOLOX官方提供的代码进行训练,训练时,不启动优先使用显存和混合精度训练,data_num_workers设置为0,batch size设置为16,以YOLOX官方提供的YOLOX-Nano模型文件作为预训练权重文件,训练300代,在训练时,前285代每10代记录一次训练信息,后15代每代记录一次训练信息,其他训练参数默认。作为对比,food-101的披萨数据集也采用以上所述方式进行训练。
为了评价YOLOX-Nano模型对菜品的检测效果,对样本进行测试后计算AP50、AP50_95(AP50和AP50_95是MS COCO比赛的map计算方式,其中AP指的是准确率(Precision)和召回率(Recall)为X,Y轴作图围成的面积,50指的是IoU设置为0.5时的AP值,50_95 指的是IoU设置为 0.5~0.95,以0.05为步长计算AP的平均值)以及total_loss来对模型进行评价。
五、结果与分析
训练结果与分析
在完全相同的训练条件下,无论是采集的数据集还是food-101的披萨数据集,总损失值都有一定的降低,模型在采集的数据集和food-101的披萨数据集上的AP50在训练300代之后分别为99.4%和90.8%,AP50_95分别为71.9%和79.8%。从训练的效果可以看出,无论是采集的数据集还是food-101的披萨数据集,AP50和AP50_95在训练第10代之后都达到了较高值,这也说明了YOLOX采用解耦头之后,收敛速度快的特点。