简明扼要理解YOLOX

由于疫情已经居家办公2周了，明天就可以正常通勤上班了，内心还有点小小的期待呢。趁着这三月暖暖的春风和屋外喳喳的鸟叫声，咱们来借鉴着大神的文章梳理一下YOLOX算法。如果对YOLO 系列算法还不了解的同学，请先找找相关的文章了解一下也可以看看我前面的这篇：简明扼要理解YOLO v3https://blog.csdn.net/rh8866/article/details/123620540

首先来看看YOLOX的网络结构图

如上图将这个网络分为四个部分：1，输入端；2，Backbone；3，Neck；4 Prediction

输入端

在yolox的输入端采用了 Mosaic、Mixup两种数据增强方式，采用了这两种数据增强，直接将Yolov3 baseline，提升了2.4个百分点；可见数据增强对网络还是很有贡献的。

（1） Mosaic数据增强

Mosaic增强的方式，是U版YOLOv3引入的一种非常有效的增强策略。而且在Yolov4、Yolov5算法中，也得到了广泛的应用。通过随机缩放、随机裁剪、随机排布的方式进行拼接，对于小目标的检测效果提升，还是很不错的。在网上找到一个很形象的图示：

（2） MixUp数据增强

MixUp是在Mosaic基础上，增加的一种额外的增强策略。

其实方式很简单，比如我们在做人脸检测的任务。

先读取一张图片，图像两侧填充，缩放到640*640大小，即Image_1，人脸检测框为红色框。

再随机选取一张图片，图像上下填充，也缩放到640*640大小，即Image_2，人脸检测框为蓝色框。

然后设置一个融合系数，比如上图中，设置为0.5，将Image_1和Image_2，加权融合，最终得到右面的Image。

从右边得图可以看出，人脸的红色框和蓝色框是叠加存在的。

我们知道，在Mosaic和Mixup的基础上，Yolov3 baseline增加了2.4个百分点。

不过有两点需要知道：

（1）在训练的最后15个epoch，这两个数据增强会被关闭掉。

而在此之前，Mosaic和Mixup数据增强，都是打开的，这个细节需要注意。

（2）由于采取了更强的数据增强方式，作者在研究中发现，ImageNet预训练将毫无意义，因此，所有的模型，均是从头开始训练的。

当然YOLOX 中也可以不使用这两个数据增强，我个人在项目使用中目前还没有使用数据增强这两个功能。

Backbone

Yolox-Darknet53的Backbone主干网络，和原本的Yolov3 baseline的主干网络都是一样的。都是采用Darknet53的网络结构，

其他得没有要特别说的，这里说一下SPP结构

SPP结构通过pooling将不同尺度的特征融合到一起，实现了数据的多尺度输入

Neck

在Neck结构中，Yolox-Darknet53和Yolov3 baseline的Neck结构，也是一样的，都是采用FPN的结构进行融合。

如上图所示，FPN自顶向下，将高层的特征信息通过上采样的方式进行传递融合，得到进行预测的特征图。而在Yolov4、Yolov5、甚至后面的Yolox-s、l等版本中，都是采用FPN+PAN的形式。

Prediction

输出层是我们要着重理解的部分，主要从四个方面进行讲解：Decoupled Head、Anchor Free、标签分配、Loss计算。

（1）Decoupled Head

在Yolox中，作者增加了三个Decoupled Head，俗称“解耦头”

首先为什么要使用解耦头呢，

作者在改进Yolov3 baseline的时候，将Yolov3 baseline 中Yolo Head，也修改为Decoupled Head。发现AP值，从38.5，增加到39.6。

不经如此作者在实验中还发现，不单单是精度上的提高。替换为Decoupled Head后，网络的收敛速度也加快了。

因此可以得到一个非常关键的结论：

目前Yolo系列使用的检测头，表达能力可能有所欠缺，没有Decoupled Head的表达能力更好。

这里看看Yolo Head和Decoupled Head的对比曲线：

有上图曲线表明：Decoupled Head的收敛速度更快，且精度更高一些。

但是需要注意的是：将检测头解耦，会增加运算的复杂度。

因此作者经过速度和性能上的权衡，最终使用 1个1x1 的卷积先进行降维，并在后面两个分支里各使用了 2个3x3 卷积，最终调整到仅仅增加一点点的网络参数。

而且这里解耦后，还有一个更深层次的重要性：

Yolox的网络架构，可以和很多算法任务，进行一体化结合。

比如：

（1）YOLOX + Yolact/CondInst/SOLO ，实现端侧的实例分割。

（2）YOLOX + 34 层输出，实现端侧人体的 17 个关键点检测。

（2） Decoupled Head的细节

我们取出其中一个检测头看看细节

从图上可以看出，Concat前总共有三个分支：

（1）cls_output：主要对目标框的类别，预测分数。因为COCO数据集总共有80个类别，且主要是N个二分类判断，因此经过Sigmoid激活函数处理后，变为202080大小。

（2）obj_output：主要判断目标框是前景还是背景，因此经过Sigmoid处理好，变为20201大小。

（3）reg_output：主要对目标框的坐标信息（x，y，w，h）进行预测，因此大小为20204。

最后三个output，经过Concat融合到一起，得到202085的特征信息。

当然，这只是Decoupled Head①的信息，再对Decoupled Head②和③同样的进行处理

Decoupled Head②输出特征信息，并进行Concate，得到404085特征信息。

Decoupled Head③输出特征信息，并进行Concate，得到808085特征信息。

再对①②③三个信息，进行Reshape操作，并进行总体的Concat，得到8400*85的预测信息。

并经过一次Transpose，变为85*8400大小的二维向量信息。

这里的8400，指的是预测框的数量，而85是每个预测框的信息（reg，obj，cls）。

有了预测框的信息，下面我们再了解，如何将这些预测框和标注的框，即groundtruth进行关联，从而计算Loss函数，更新网络参数呢？

（2）Anchor Free

目前行业内，主要有Anchor Based和Anchor Free两种方式。

在Yolov3、Yolov4、Yolov5中，通常都是采用Anchor Based的方式，来提取目标框，进而和标注的groundtruth进行比对，判断两者的差距。

① Anchor Based方式

比如输入图像，经过Backbone、Neck层，最终将特征信息，传送到输出的Feature Map中。

这时，就要设置一些Anchor规则，将预测框和标注框进行关联。从而在训练中，计算两者的差距，即损失函数，再更新网络参数。

比如在下图的，最后的三个Feature Map上，基于每个单元格，都有三个不同尺寸大小的锚框。

以输入图像大小416*416为例：

当输入为416416时，网络最后的三个特征图大小分别为1313，2626，5252。

我们可以看到，黄色框为小狗的Groundtruth，即标注框（真值）。

而蓝色的框，为小狗中心点所在的单元格（红框），所对应的锚框，每个单元格都有3个蓝框。

当采用COCO数据集，即有80个类别时。

基于每个锚框，都有x、y、w、h、obj（前景背景）、class（80个类别），共85个参数。

因此会产生3*(1313+2626+5252）85=904995个预测结果。

如果将输入从416416，变为640640，最后的三个特征图大小为2020,4040,80*80。

则会产生3（2020+4040+8080）*85=2142000个预测结果。

② Anchor Free方式

在前面Anchor Based中，每个Feature map的单元格，都有3个大小不一的锚框。

那么Yolox-Darknet53就没有吗？

其实并不然，这里只是巧妙的，将前面Backbone中，下采样的大小信息引入进来。

比如上图中，最上面的分支，下采样了5次，2的5次方为32。

并且Decoupled Head①的输出，为202085大小。

因此如上图所示：

最后8400个预测框中，其中有400个框(2020)，所对应锚框的大小，为3232。

同样的原理，中间的分支，最后有1600个预测框(4040)，所对应锚框的大小，为1616。

最下面的分支，最后有6400个预测框(8080)，所对应锚框的大小，为88。

当有了8400个预测框的信息，每张图片也有标注的目标框的信息。

这时的锚框，就相当于桥梁。

这时需要做的，就是将8400个锚框，和图片上所有的目标框进行关联，挑选出正样本锚框。

而相应的，正样本锚框所对应的位置，就可以将正样本预测框，挑选出来。

这里采用的关联方式，就是标签分配。

（3）标签分配

标签分配这部分我觉得是需要着重讲一下。

当有了8400个Anchor锚框后，这里的每一个锚框，都对应85*8400特征向量中的预测框信息。

不过需要知道，这些预测框只有少部分是正样本，绝大多数是负样本。

那么到底哪些才是正样本呢？

这里需要利用锚框和实际目标框的关系，先挑选出一部分适合的正样本锚框。

比如第3、10、15个锚框是正样本锚框，则对应到网络输出的8400个预测框中，第3、10、15个预测框，就是相应的正样本预测框。

训练过程中，在锚框的基础上，不断的预测，然后不断的迭代，从而更新网络参数，让网络损失越来越小，预测的越来越准。

那么在Yolox中，是如何挑选正样本锚框的呢？

这里就涉及到两个关键点：初步筛选、SimOTA。

① 初步筛选

初步筛选的方式主要有两种：根据中心点来判断、根据目标框来判断；

这部分的代码，在models/yolo_head.py的get_in_boxes_info函数中。

a） 根据中心点来判断：

规则：寻找anchor_box中心点，落在groundtruth_boxes矩形范围的所有anchors。

通过groundtruth的[x_center,y_center，w，h]，计算出每张图片的每个groundtruth的左上角、右下角坐标，左上角（gt_l，gt_t），右下角（gt_r，gt_b）。

如下图：

groundtruth的矩形框范围确定了，再根据范围去选择适合的锚框。这里再绘制一个锚框的中心点，（x_center，y_center），*(这里不是groundtruth的（x_center,y_center）)*。

而右面的图片，就是寻找锚框和groundtruth的对应关系。

即计算锚框中心点（x_center，y_center），和人脸标注框左上角（gt_l，gt_t），右下角（gt_r，gt_b）两个角点的相应距离。通过计算将中心点落在groundtruth矩形框内的锚框挑选出来。

b）根据目标框来判断：

除了根据锚框中心点，和groundtruth两边距离判断的方式外，作者还设置了根据目标框判断的方法。

规则：以groundtruth中心点为基准，设置边长为5的正方形，挑选在正方形内的所有锚框。

将落中心点落在以groundtruth中心点边长为5的正方形范围内的所有的anchors，都提取出来。

经过上面两种挑选的方式，就完成初步筛选了，挑选出一部分候选的anchor，进入下一步的精细化筛选。

② 精细化筛选

而在精细化筛选中，就用到论文中提到的SimOTA了：

从提升效果上来看，引入SimOTA后，AP值提升了2.3个百分点，还是非常有效的。

我们将SimOTA的前后流程进行拆解，看一下是如何进行精细化筛选的？

（如下部分是yolox之前版本没有使用的，我认为这部分也是YOLOX的精华所在，值得理解透彻）

整个筛选流程，主要分为四个阶段：

a.初筛正样本信息提取

b.Loss函数计算

c.cost成本计算

d.SimOTA求解

接下来为了便于理解，我们假定图片上有3个目标框，即3个groundtruth。

再假定目前在做的项目是对人脸和人体检测，因此检测类别是2。

通过前面介绍，我们知道有8400个锚框，但是经过初步筛选后，假定有1000个锚框是正样本锚框。

a.初筛正样本信息提取

初筛出的1000个正样本锚框的位置，我们是知道的。

而所有锚框的位置，和网络最后输出的85*8400特征向量是一一对应。

所以根据位置，可以将网络预测的候选检测框位置bboxes_preds、前景背景目标分数obj_preds、类别分数cls_preds等信息，提取出来。

以前面的假定信息为例，

bboxes_preds是候选检测框的信息，因此维度为[1000，4]。

obj_preds因为是目标分数，所以维度是[1000，1]。

cls_preds因为是类别分数，所以维度是[1000，2]。

b.Loss函数计算

针对筛选出的1000个候选检测框，和3个groundtruth计算Loss函数。

可以计算出3个目标框，和1000个候选框，每个框相互之间的iou信息，再通过-torch.log计算，得到位置损失；

将类别的条件概率和目标的先验概率做乘积，得到目标的类别分数。

再通过F.binary_cross_entroy的处理，得到3个目标框和1000个候选框的综合loss值。

c.cost成本计算

有了reg_loss和cls_loss，就可以将两个损失函数加权相加，计算cost成本函数了。

这里涉及到论文中提到的一个公式：

相应的，对应于yolo_head.py的get_assignments函数中的代码：

可以看出，公式中的加权系数，即代码中的3.0

**d.SimOTA **

有了上面的一系列信息，标签分配问题，就转换为了标准的OTA问题。

但是经典的Sinkhorn-Knopp算法，需要多次迭代求得最优解。

作者也提到，该算法会导致25%额外训练时间，所以采用一种简化版的SimOTA方法，求解近似最优解。这里对应的函数，是get_assignments函数中的self.dynamic_k_matching：

其中的流程如下：

第一步：设置候选框数量

首先按照cost值的大小，新建一个全0变量matching_matrix，这里是[3,1000]。

通过上面第二行代码，设置候选框数量为10。

再通过第三行代码，从前面的pair_wise_ious中，给每个目标框，挑选10个iou最大的候选框。

因为前面假定有3个目标，因此这里topk_ious的维度为[3，10]。

第二步：通过cost挑选候选框

下面再通过topk_ious的信息，动态选择候选框，这里是个关键。

代码如dynamic_k_matching函数中，下图所示：

为了便于大家理解，大白先把第一行制作成图示效果。

这里的topk_ious，是3个目标框和预测框中，最大iou的10个候选框：

经过torch.clamp函数，得到最终右面的dynamic_ks值。

我们就知道，目标框1和3，给他分配3个候选框，而目标框2，给它分配4个候选框。

那么基于什么标准分配呢？

这时就要利用前面计算的cost值，即[3,1000]的损失函数加权信息。

在for循环中，针对每个目标框挑选相应的cost值最低的一些候选框。

比如右面的matching_matrix中，cost值最低的一些位置，数值为1，其余位置都为0。

因为目标框1和3，dynamic_ks值都为3，因此matching_matrix的第一行和第三行，有3个1。

而目标框2，dynamic_ks值为4，因此matching_matrix的第二行，有4个1。

第三步：过滤共用的候选框

不过在分析matching_matrix时，我们发现，第5列有两个1。

这也就说明，第五列所对应的候选框，被目标检测框1和2，都进行关联(一个锚框只能检测一个目标)。

因此对这两个位置，还要使用cost值进行对比，选择cost较小的值，再进一步筛选。

这里为了便于理解，还是采用图示的方式：

首先第一行代码，将matching_matrix，对每一列进行相加。

这时anchor_matching_gt中，只要有大于1的，说明有共用的情况。

上图案例中，表明第5列存在共用的情况。

再利用第三行代码，将cost中，第5列的值取出，并进行比较，计算最小值所对应的行数，以及分数。

我们将第5列两个位置，假设为0.4和0.3。

经过第三行代码，可以找到最小的值是0.3，即cost_min为0.3，所对应的行数，cost_argmin为2。

经过第四行代码，将matching_matrix第5列都置0。

再利用第五行代码，将matching_matrix第2行，第5列的位置变为1。

最终我们可以得到3个目标框，最合适的一些候选框，即matching_matrix中，所有1所对应的位置。

(4）Loss计算

经过第三部分的标签分配，就可以将目标框和正样本预测框对应起来了。

下面就可以计算两者的误差，即Loss函数。

计算的代码，位于yolo_head.py的get_losses函数中。

我们可以看到：

检测框位置的iou_loss，Yolox中使用传统的iou_loss，和giou_loss两种，可以进行选择。

而obj_loss和cls_loss，都是采用BCE_loss的方式。

当然除此之外，还有两点需要注意：

a.在前面精细化筛选中，使用了reg_loss和cls_loss，筛选出和目标框所对应的预测框。

因此这里的iou_loss和cls_loss，只针对目标框和筛选出的正样本预测框进行计算。

而obj_loss，则还是针对8400个预测框。

b.在Decoupled Head中，cls_output和obj_output使用了sigmoid函数进行归一化，

但是在训练时，并没有使用sigmoid函数，原因是训练时用的nn.BCEWithLogitsLoss函数，已经包含了sigmoid操作。

而在推理过程中，是使用Sigmoid函数的。

以上便是对YOLOX的重点部分的说明，当然也是踩巨人肩膀上完成的，文中的图示很多都是从其他地方拿过来的。以上文中如有错误的部分或者没有说明白的欢迎留言一起交流。