论文阅读|目标检测之PSS：更简单有效的End-to-End检测

paper：arxiv.org/abs/2101.11782

code：txdet/FCOSPss (github.com)

摘要

PSS是一种简单的NMS-free、end-to-end的目标检测框架，仅需要对现有one-stage检测器（FCOS、ATSS）做很小的修改，即可移除复杂的后处理NMS，达到甚至超过原有检测器的性能，而且几乎不影响推理速度。这里实现NMS-free的方式是给每个gt只激活一个正样本，为此，需要在现有NMS-based的one-stage detector基础上引入一个正样本选择器（positive sample selector，PSS），作为单独一个分支来辅助完成为每个object挑选最优正样本的工作，基础detector以FCOS为例，具体网络结构见图1 。由于这里的学习目标（损失函数）同时涉及一对多和一对一的标签分配，所以有一些训练样本的标签存在冲突，使得模型难以训练，产生优化冲突。于是作者通过引入 stop-grad解决了这个问题。 该算法在COCO数据集上比使用NMS的FCOS baseline以及很多最近的NMS-free的检测器都好。

1. Introduction

自从anchor-free的出现，大家意思到了anchor boxes不是必须的，因此省去了很多关于anchor的人为设置，整个detector的训练pipeline中只剩下NMS后处理需要认为设置。而NMS之所以还在很多检测器中保留则是因为，在很多检测器训练过程过程中，一个gt物体往往会有多个对应的正样本，因此在推断时需要采用一种机制（NMS）来从这些正样本边框中选择最好的一个。

最近的DeTR等基于transformer的检测器将目标检测问题转换成一个set-to-set的预测问题，利用了匈牙利匹配算法来为每个gt box选择一个正样本。匈牙利算法会自动的为每个gt选择一个最佳的正样本，使得模型输出的损失最小，从而可以移除NMS，变成一个完全的端到端（end-to-end）的可训练模型。DeFCN或者OneNet等非基于transformer的全卷积网络也是先了端到端的NMS-free目标检测。所以本篇论文的工作要继续沿着NMS-free这条线，设计一个简单、高性能、NMS-free、完全全卷积的，可端到端训练的目标检测器。作者在FCOS的基础上做修改，如图1所示，在FOCS的head上增加了一个**“positive sample selector（PSS）”head分支用于为每个gt选择最佳的正样本，同时作者还重新设计的损失函数（目标函数）**。

对于网络的训练，保持了FCOS原来的分类损失，而上文提到的标签分配差异导致的训练冲突通过stop-grad操作来停止PSShead的梯度传播到原始的FCOS网络参数，见图1. stop-grad在这样的情况下非常有用：当网络包括两个子网络A和B，其中B的优化依赖于A的收敛，而A的优化却不太依赖于B，从而是非对称。图3也展示了stop-grad的有效性。

该方法不依赖于注意力机制，如Transformer，或者跨尺度特征聚合，如3D max filtering。

本文提出的检测器称作$FCO{S}_{PSS}$，具有以下优点：

消除了NMS，所以更加简单。基于FCOS，所以继承了FCOS简单性，能够应用到任何FCN可解决的任务。

和原版FCOS相比，引入一个紧凑的PSS head来消除了NMS，相比原版FCOS，带来的额外计算可忽略。

灵活。一旦$FCO{S}_{PSS}$训练完成，可以选择丢弃PSS head，当作标准FCOS来使用。

精度的提高。

能够应用到其它基于anchor box的检测器，比如RetinaNet。我们将PSS head添加到修改后的RetinaNet（每个位置一个anchor，并且应用ATSS的标签分配策略）取得了很好的结果。

该思想能够适用于其它实力识别的任务，比如能让基于FCOS的实例分割（SOLO），关键点检测以及目标追踪（SiamCAR）等。

2. Related Work

DeFCN首次证明不需要依赖于self-attention机制的sequence-to-sequence（set-to-set）学习也可以实现NMS-free，而OneNet于DeFCN相似，采用一种one-to-one的标签分配方式，并且加了cascade head来帮助训练。但是OneNet的性能不及DeFCN。作者受DeFCN启发，设计了一个一对一标签分配策略，如等式（3）所示。显然该标签分配策略最终NMS-free检测精度有直接影响。我们和DeFCN的主要区别在于我们使用PSS head来实现给每个gt选择一个正样本，而DeFCN则是通过一个3D max filtering来实现的，3D max filtering聚合多层特征来抑制重复检测。还有个区别在于本工作尽可能地保持FCOS的原样。

3. Our Method

网络的整体架构见图1，很清晰，相比FCOS，唯一的修改就是增加了一个有两个额外卷积层的PSS head。下面分析整体的目标函数（训练目标或者说损失函数）

3.1 Overall Training Objective

总的训练目标如下：
$$\mathcal{L}={\mathcal{L}}_{fcos}+{\lambda }_1\cdot {\mathcal{L}}_{pss}+{\lambda }_2\cdot {\mathcal{L}}_{rank}.(1)$$
其中${\lambda }_1,{\lambda }_2$是平衡系数，${\mathcal{L}}_{fcos}$包含和原来FCOS中一样的损失项，即分类loss（focal loss）、回归loss（GIoU loss）和center-ness loss。ranking loss不核心部分，所以$GIoU loss

其中${\lambda }_1,{\lambda }_2$是平衡系数，${\mathcal{L}}_{fcos}$包含和原来FCOS中一样的损失项，即分类loss（focal loss）、回归loss（GIoU loss）和center-ness loss。ranking loss不核心部分，所以${\lambda }_2$设置为0.25.表9显示其仅仅提高轻微最终检查精度，约0.2-0.3的mAP，而之所以将它加进来是因为不会带来太多复杂度。

3.1.1 PSS Loss Lpss

${\mathcal{L}}_{pss}$是作用于PSS head的分类loss，是训练PSS head从而完成one-to-one pred的关键loss，在one-to-many pred中一个object往往对应多个正样本，PSS head就是用来从多个正样本中挑选最优的那一个作为唯一正样本。这里假设一对一正标签分配能够保证每个gt有且只有一个正样本。如图1所示，PSS head的输出是一个大小为$R^{H\times W\times 1}$ binary mask，假设$\sigma (pss)$是图上一个点，如果它对应于一个实例的正样本，则其学习对象为1，反之为负标签，那么最简单是就是设计一个二分类器，然后利用FCOS的C通道的分类器来得到一个用于C通道分类器的损失项。

具体而言，该损失项计算的是$\sigma (pss)\cdot \sigma (s)\cdot \sigma (ctr)$和gt之间的focal loss，其中$\sigma$为sigmoid function， s为原始detector的分类score， ctr为center-ness。通过这种方式训练，使得PSS有能力为每个object激活一个最优的正样本。

3.1.2 Ranking Loss Lrank

ranking loss能给NMS-free的检测器带来精度提高，所以给每个训练图片增加等式2的损失函数。
$${\mathcal{L}}_{\text{rank }}=\frac{1}{n_{-}{n}_{+}}\sum _{i_{-}}^{n_{-}}\sum _{i_{+}}^{n_{+}}\max \left(0,\gamma -{\hat{P}}_{i_{+}}\left(c_{i_{+}}\right)+{\hat{P}}_{i_{-}}\left(c_{i_{-}}\right)\right).(2)$$
这里γ是表示正负样本见margin的超参数（使得二者建有一个区分带？），实验中默认γ=0.5，$n_{\_},n_{+}$分别指代正负样本数，${\hat{P}}_{i_{+}}\left(c_{i_{+}}\right)$指的是属于类别$c_{i+}$的正样本$i_{+}$的分类得分，${\hat{P}}_{i_{-}}\left(c_{i_{-}}\right)$指的是属于类别$c_{i-}$的正样本$i_{-}$的分类得分，实验中，我们从所有负样本中选择${\hat{P}}_{i_{-}}\left(c_{i_{-}}\right)$得分前$n_{-}$的，$n_{-}$通常取100.

3.2. One-to-many Label Assignment

一对多的标签分配被广泛用于目标检测任务。因为对于目标检测任务而言，gt 边框的标注存在模糊性：即使偏离一些像素，但是因为仍然包括了该实例，所以也被视为正训练样本。这也是采用静态规则的一对一分配难以取得满意结果的原因，因为我们很难找到一个很好定义的注释。

一对多分配的优势在于每个gt有多个box，能够学习得到一个强大的分类器，具有比例和变换不变性等。但是，其结果就是每个gt周围有多个检查框，使得NMS不可避免。虽然需要NMS，但是一对多标签分配有其优势：1）有丰富的训练样本帮助学习更具不变性的检测器。2）与大多数的数据增强方法保持一致。所以，作者相信保持FCOS训练对象与一对多标签分配的联系是关键的。

3.3. One-to-one Label Assignment

当进行one-to-one label assign的时候，最重要的就是给每个gt j（$c_j,b_j$表示gt的类别标签和边框坐标）挑选最优的那个正样本anchor i（这里anchor表示anchor box或者anchor point），那就需要给每个样本进行质量评估，DeFCN指出最优的匹配的时候应该同时考虑classification的匹配和localization的匹配，借鉴DeFCN，我们对样本评估matching score $Q_{i,j}$定义如下：
$$Q_{i,j}=\underset{\text{positiveness prior }}{\underbrace{1\left[i\in {\Omega }_j\right]}}\cdot \underset{\text{classification }}{\underbrace{{\left[{\hat{P}}_i\left(c_j\right)\right]}^{1-\alpha }}}\cdot \underset{\text{localization }}{\underbrace{{\left[\mathrm{IoU}\left({\hat{b}}_i,b_j\right)\right]}^{\alpha }}}.（3）$$
其中
$${\hat{P}}_i\left(c_j\right)=\sigma \left({\mathrm{p}\mathrm{s}\mathrm{s}}_i\right)\cdot \sigma \left(s_i\right)\cdot \sigma \left({\mathrm{ctr}}_i\right).（4）$$
这里$s_i,ct{r}_i$指的是anchor i的分类得分和center-ness的预测。$ps{s}_i$表示一个二值掩膜预测得分（binary mask prediction scores），也即PSS head的输出。${\hat{P}}_i(c_j)$指的是anchor i属于类别$c_j$的分类得分，sigmoid函数$\sigma$旨在将得分归一化成一个概率。假设等式（4）中的三个概率独立，所以它们相乘得到${\hat{P}}_i(c_j)$。 ${\hat{b}}_i,b_j$分别表示预测的anchor i的边框坐标和gt j的边框坐标。超参数$\alpha \in [0,1]$用于调整分类和定位比例。

并不是所有的anchor都能分配为正样本，尤其是处于gt边框外的anchor。这里${\Omega }_j$表示实例j的候选正样本anchor的集合。FCOS限制只有处于$b_j$中心的anchor points才属于${\Omega }_j$，RetinaNet则是通过anchors和$b_j$的IoU限制得到${\Omega }_j$。${\Omega }_j$的设计会极大地影响model的性能。

在本算法中，${\Omega }_j$简单地采用原来detector中使用的正样本，如等式3所示。比如$FCO{S}_{PSS}$和FCOS一样使用实例j中心的anchor points作为${\Omega }_j$，而$ATS{S}_{PSS}$则是使用ATSS中的采样策略。

最后，通过DeTR或者DeFCN中的匈牙利算法来接二分图匹配即可为一张图片中的每个gt分配一个label，即通过给每个实例j寻找最佳的anchor index i来最大化${\sum }_j{Q}_{i,j}$。当然，使用简单的top1选择来替代匈牙利算法也能取得差不多的性能。

3.4. Conflict in the Two Classification Loss Terms

回顾等式（1）可知，我们需要最小化两个分类损失，一个是原始FCOS中的一对多标签的分类损失${\mathcal{L}}_{fcos}$，这时每个gt可以被分配k个正样本。第二个是PSS的${\mathcal{L}}_{pss}$, 其主要责任是从剩下的样本（第一次分配的k个正样本中）中选择一个正样本。所以，训练PSS分类器时，${\mathcal{L}}_{fcos}$的k个正样本中的剩下k-1个会被分配为负标签。模型很难同时拟合这个两个目标函数。也就是说训练过程中多个样本可能同时被当作正样本和负样本，为了解决这个问题，我们引入 stop-grad。使得PSS head的梯度（图1中的虚线部分）不会传回原始的FCOS网络（图1中除了虚线的部分）中，使得PSS head对原来FCOS的训练影响降到最低（但是因为${\mathcal{L}}_{pss}$与${\mathcal{L}}_{fcos}$存在耦合，所以还是有影响）。

另一种表述：

one-to-many的分类loss和one-to-one的分类loss是存在优化矛盾的。比如一个gt在one-to-many里对应K个正样本，但通过我们的PSS head仅选出了一个最优的正样本，也就是需要把其余K-1个样本置为负样本，也就是说多个样本可能同时被当作正样本和负样本，这使得模型难以训练，产生优化冲突。为了解决这个问题，我们引入 stop-grad，如框架图所示，作用于PSS head上，也是受到了Kaiming的那篇《Exploring Simple Siamese Representation Learning》启发。

3.5. Stop Gradient

stop-gradient操作从数学上看就是在网络训练时让网络一部分为常数。具体到这里，则是在SGD更新原来FCOS参数的时候让PSS head 的参数为常量，使得没有梯度从PSS head传回网络剩余部分，这样解决优化冲突。

设网络要优化的参数分为两部分$\theta ={\theta }_{fcos},\theta pss$，本质上来讲stop-gradient相当于交替的优化这两个集合的参数。

即，我们想要求解：
$$\underset{{\theta }_{fcos},{\theta }_{pss}}{\min }\mathcal{L}({\theta }_{fcos},{\theta }_{pss}).$$
我们可以交替的解决这两个子问题（t表示iterations）：
$${\mathbf{\theta }}_{fcos}^t\leftarrow \mathrm{argmin}{\mathbf{\theta }}_{fcos}\mathcal{L}\left({\mathbf{\theta }}_{fcos},{\mathbf{\theta }}_{pss}^{t-1}\right);(5)$$
和,
$${\mathbf{\theta }}_{pss}^t\leftarrow {\mathrm{argmin}}_{{\mathbf{\theta }}_{pss}}\mathcal{L}\left({\mathbf{\theta }}_{fcos}^t,{\mathbf{\theta }}_{pss}\right)。（6）$$
当求解等式5时，相对于pss的参数${\theta }_{pss}$为0.

我们这里使用的带有stop-gradient 的SGD近似于上述的交替优化，但是实际上我们并非单独的优化这两个子问题。

当然，也可以在一个iteration中解上述两个子问题，即先让${\theta }_{pss}$初始化为0，然后求解等式5直到收敛，然后在求解等式6知道收敛，这让就等同于训练原始的FCOS到收敛，然后freezeFCOS，再训练PSS head直到收敛。而我们的实验表明这种方式会导致较次的精度，而且训练时间更长。

我们通过*detach()*操作（stop-gradient）来断绝PSS head的优化尽量不影响前面的特征，实现解耦，并且还能利用上前面收敛较快的鲁棒的reg 特征。我们实验发现这个操作能随着训练持续提高one-to-one pred (w/ NMS)的检测性能。如图3：

4.Experiments

4.1 Implementation Details

模型基于MMDetection。使用带有NMS后处理的检测器FCOS和ATSS作为基准检测器，然后加上PSS head后消除后处理NMS。

细节如下

All the ablation comparisons are based on the ResNet50 [8] backbone with FPN [12], and the feature weights are initialized by the pretrained ImageNet model. Unless otherwise specified, we train all the models with the ‘3×’ training schedule (36 epochs). Specifically, we rain the models using SGD on 8 Tesla-V100 GPUs, with an initial learning rate of 0.01, a momentum of 0.9, a weight decay of 10−4, a mini-batch size of 16. The learning rate decays by a factor of 10 at the 24th and 33th epoch respectively

4.2. Ablation Studies

我们主要在FCOS和ATSS上接上我们的PSS方法进行实验，主要结果如表1：

我们大大缩小了E2E检测与NMS-based检测的gap，达到甚至超过了ATSS、FCOS的baseline，网络推理耗时仅少量增加，并且由于移除了NMS，我们的后处理耗时减小了。

不同方法的p5-p7 level的分类得分的可视化如下：

Stop Gradient vs. Two-step Training

这里的two-step指的是“first train the FCOS/ATSS model to convergence, and then train the PSS module alone by freezing the FCOS/ATSS model”。如前面提到的，two-step的结果较次。

Attaching PSS to Regression vs. Classification Branch

测试PSS head放在不同分支，结果表明放在reg 分支表现最好，所以所有实验都放reg分支。

How Many Conv. Layers for the PSS Head

结果表明，两个conv最好。

Effect of the Center-ness Branch

加了能提高一点精度。

**Loss Weight of PSS Loss $L_{pss}\ast \ast$

等式1中$L_{pss}$的参数${\lambda }_1$可以取值{0.5，1，2}，结果表明取值1最好。

Matching Score Function

等式3里，匹配得分函数的两部分可以是相加或者相乘，原来是相乘，这里消融实验尝试了相加，改为相加后的形式为：
$$(1-\alpha )\cdot Pˆi(cj)+\alpha \cdot IoU(ˆbi,bj).$$
结果表明还是相乘好，相乘的时候$\alpha$取值0.8表现最好。

Effect of the Ranking Loss

加了rank loss能提高些精度~0.2。

5.Conclusion

我们提出了一个更简单和更有效的E2E检测框架，仅需要对FCOS、ATSS进行简单的修改，就能移除NMS，并且达到和超过ATSS、FCOS的baseline。因为我们保留了原始detector，所以模型训练完成后，还可以继续选择使用NMS的结果，而去除NMS的pipeline使得我们模型更加容易部署。我们希望我们的工作能够带给大家一些关于End-to-End检测的新的启发。

本质上就是再原模型样本分配后的正样本基础上再做一对一分配，而这时候的一对一分配方式可以是匈牙利算法也可以是top1算法，为此增加了一个head来专门预测一个二值掩膜来表示第二次一对一分配中有对应gt的那些，当然这样之后的修改loss函数，其实就是再FCOS的基础上增加一个pss loss，这个loss就是二值mask的输出和FCOS原来的分类输出和cente-ness输出相乘后与gt的focal loss，然后为了解决两个分配的冲突，采用了一种stop-gradient的操作，按照作者描述，其实就是使用了一个detach（）做了分离。

启发就是loss的计算部分，相加可以改相乘，反过来也可以，然后标签分配也能存在two-stage形式的分配方式，本质上讲相当于把two-stage 的proposal预测转移到了two-stage 的标签分配。

参考

作者在知乎上的帖子：PSS：更简单有效的End-to-End检测 - 知乎 (zhihu.com)