PSS：更简单有效的End-to-End检测

简单介绍一下我们最近的一项关于end-to-end检测方面的工作：

Object Detection Made Simpler by Eliminating Heuristic NMS

论文链接：https://arxiv.org/pdf/2101.11782.pdf

代码链接：即将开源（https://github.com/txdet/FCOSPss）

1. Motivation

我们提出了一种简单的NMS-free、end-to-end的目标检测框架，仅需要对现有one-stage检测器（FCOS、ATSS）做很小的修改，即可移除复杂的后处理NMS，达到甚至超过原有检测器的性能，并且几乎不影响推理速度。试想一下，如果模型有能力为每个object仅激活一个正样本，那复杂的NMS是不是就可以移除了。为此，我们在现有NMS-based的one-stage detector基础上引入一个正样本选择器（positive sample selector，PSS），作为单独一个分支来辅助完成为每个object挑选最优正样本的工作，基础detector以FCOS为例，具体结构如下：

目标检测之前主要分为two-stage和one-stage，尤其是one-stage方法越来越受欢迎，基本是工业界做项目尤其是端上项目的首选。FCOS和FoveaBox等方法掀起了去anchor这种人工先验的潮流。那么现在对于实现真正的end-to-end检测只剩下NMS这个启发式的复杂的后处理操作了。最近，DETR、Deformable DETR等方法通过使用Hungarian matching方法实现one-to-one的匹配，提供了基于Transformer的attention机制在E2E检测方面的新思路。前段时间，培泽和剑锋分别提出了OneNet和DeFCN，不使用Transformer进行NMF-free的E2E检测，我们的方法也是主要受DeFCN的启发并进行改进和重新设计。总的来说，我们设计了一个更简单性能更强的E2E检测框架，基于现有NMS-based one-stage的检测器，引入PSS head并重新设计了目标函数。

2.Method

2.1 Overall Training Objective

整体损失函数包括三部分，原始detector的loss（这里以FCOS为例），PSS loss和ranking loss。我们尽可能不破坏原有NMS-based的detector，原有的loss也都保留，你可以认为是作为辅助loss，以FCOS为例，就是分类loss、回归loss和center-ness loss。DeFCN里面没用box loss作为aux loss，我们通过改进也用上可以work。

2.1.1 PSS loss 

是作用于PSS head的分类loss，是训练PSS head从而完成one-to-one pred的关键loss，在one-to-many pred中一个object往往对应多个正样本，PSS head就是用来从多个正样本中挑选最优的那一个作为唯一正样本。你可以认为它的输出是一个大小为  binary mask，计算的是  和gt之间的focal loss，其中  为sigmoid function，  为原始detector的分类score，  为center-ness。通过这种方式训练，使得PSS有能力为每个object激活一个最优的正样本。

2.1.2 Ranking loss 

我们早期实验发现，预测结果中会出现FP 的max socre 高于 TP min score的现象，其实也很正常，为了进一步拉开差距，从而引入来Ranking loss：

通过在image level加margin，能进一步缓解之前的实验现象，并不影响训练效率的情况下能带来一定的性能提升。

2.2 One-to-one Label Assignment

当进行one-to-one label assign的时候，最重要的就是给每个gt挑选最优的那个正样本，那就需要给每个样本进行质量评估，DeFCN指出最优的匹配的时候应该同时考虑classification的匹配和localization的匹配，借鉴DeFCN，我们对样本评估matching score：

不同点在于，除了分类质量部分我们引入来PSS 的mask和center-ness score，我们对spatial prior进行了修改。因为我们发现原先比如DeFCN固定采用center sampling策略作为spatial prior会有问题，比如FCOS在每个level assign正样本的时候是配置了range的，再比如ATSS就不是用center sampling策略，这样会导致一个问题：one-to-one match部分的正样本可能是one-to-many match部分的负样本，从而产生优化冲突。于是我们引入positiveness prior，即我们把原来的空间先验换成了原始detector里的正样本先验。

2.3 Conflict in the Two Classification Loss Terms

其实，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》启发。

2.4 Stop Gradient

因为我们发现上述优化冲突问题，尤其是可视化特征发现，原始鲁棒的reg 特征被PSS head搞sharp了。我们通过_detach()_操作来断绝PSS head的优化尽量不影响前面的特征，实现解耦，并且还能利用上前面收敛较快的鲁棒的reg 特征。我们实验发现这个操作能随着训练持续提高one-to-one pred (w/ NMS)的检测性能。如下图：

3. Experiments

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

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

3.1 Visualization

我们做了更多详细的ablation 实验，还请参考我们的论文全文。

4. Conclusion

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