没有NMS！阿里巴巴和阿大提出PSS：更简单有效的端到端目标检测

点击上方“CVer”，选择加"星标"置顶

重磅干货，第一时间送达

本文作者：SuperHui |  来源：知乎（已授权）

https://zhuanlan.zhihu.com/p/347515623

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

Object Detection Made Simpler by Eliminating Heuristic NMS

论文链接：arxiv.org/abs/2101.11782

代码链接：即将开源

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.1 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检测的新的启发。

PSS论文下载

后台回复：PSS，即可下载上述论文PDF

CV资源下载

后台回复：CVPR2020，即可下载代码开源的论文合集

后台回复：ECCV2020，即可下载代码开源的论文合集

后台回复：YOLO，即可下载YOLOv4论文和代码

后台回复：Transformer综述，即可下载两个最新的视觉Transformer综述PDF，肝起来！

重磅！CVer-细分垂直交流群成立

扫码添加CVer助手，可申请加入CVer-细分垂直方向 微信交流群，也可申请加入CVer大群，细分方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、Transformer、PyTorch和TensorFlow等群。

一定要备注：研究方向+地点+学校/公司+昵称（如目标检测+上海+上交+卡卡），根据格式备注，才能通过且邀请进群

▲长按加微信群

▲长按关注CVer公众号

整理不易，请给CVer点赞和在看！