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

技术讨论 hello_uncle ⋅ 于 4周前 ⋅ 133 阅读

作者丨SuperHui@知乎(已授权)
来源丨https://zhuanlan.zhihu.com/p/347515623
编辑丨极市平台

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

Object Detection Made Simpler by Eliminating Heuristic NMS

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

代码链接:即将开源

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 $\mathcal{L}_{pss}$

$\mathcal{L}_{pss}$ 是作用于PSS head的分类loss,是训练PSS head从而完成one-to-one pred的关键loss,在one-to-many pred中一个object往往对应多个正样本,PSS head就是用来从多个正样本中挑选最优的那一个作为唯一正样本。你可以认为它的输出是一个大小为 $R^{H \times W \times 1}$ binary mask,计算的是 $\sigma(pss) \cdot \sigma(s) \cdot \sigma(ctr)$ 和gt之间的focal loss,其中 $\sigma$ 为sigmoid function, $s$ 为原始detector的分类score, $ctr$ 为center-ness。通过这种方式训练,使得PSS有能力为每个object激活一个最优的正样本。

2.1.1 Ranking loss $\mathcal{L}_{rank}$

我们早期实验发现,预测结果中会出现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检测的新的启发。

大叔

成为第一个点赞的人吧 :bowtie:
回复数量: 0
暂无回复~
您需要登陆以后才能留下评论!