我们基于 FCOS,首次在 dense prediction 上利用全卷积结构做到 E2E,即无 NMS 后处理。我们首先分析了常见的 dense prediction 方法(如 RetinaNet、FCOS、ATSS 等),并且认为 one-to-many 的 label assignment 是依赖 NMS 的关键。受到 DETR 的启发,我们设计了一种 prediction-aware one-to-one assignment 方法。此外,我们还提出了 3D Max Filtering 以增强 feature 在 local 区域的表征能力,并提出用 one-to-many auxiliary loss 加速收敛。我们的方法基本不修改模型结构,不需要更长的训练时间,可以基于现有 dense prediction 方法平滑过渡。我们的方法在无 NMS 的情况下,在 COCO 数据集上达到了与有 NMS 的 FCOS 相当的性能;在代表了密集场景的 CrowdHuman 数据集上,我们的方法的 recall 超越了依赖 NMS 方法的理论上限。
整体方法流程如下图所示:
One-to-many vs. one-to-one
自 anchor-free 方法出现以来,NMS 作为网络中最后一个 heuristic 环节,一直是实现 E2E dense prediction 的最大阻碍。但其实我们可以发现,从 RPN、SSD、RetinaNet 等开始,大家一直遵循着这样一个流程:先对每个目标生成多个预测(one-to-many),再将多个预测去重(many-to-one)。所以,如果不对前一步 label assignment 动刀,就必须要保留去重的环节,即便去重的方法不是 NMS,也会是 NMS 的替代物(如 RelationNet,如 CenterNet 的 max pooling)。
那直接做 one-to-one assignment 的方法是否存在呢?其实是有的。上古时代有一个方法叫MultiBox,对每个目标和每个预测做了 bipartite matching,DETR 其实就是将该方法的网络换成了 Transformer。此外还有一个大家熟知的方法:YOLO,YOLO 也是对每个目标只匹配一个 grid,只不过它是采用中心点做的匹配,而且有 ignore 区域。
Prediction-aware one-to-one
于是接下来的问题就是,在 dense prediction 上我们能不能只依赖 one-to-one label assignment,比较完美地去掉 NMS?我们首先基于去掉 centerness 分支的 FCOS,统一网络结构和训练方法,用 Focal Loss + GIoU Loss,做了如下分析实验:
我们设计了两种 hand-crafted one-to-one assignment 方法,分别模仿 RetinaNet(基于 anchor box)和 FCOS(基于 center 点),尽可能做最小改动,发现已经可以将有无 NMS 的 mAP 差距缩小到 4 个点以内。
但我们认为手工设计的 label assignment 规则会较大地影响 one-to-one 的性能,比方说 center 规则对于一个偏心的物体就不够友好,而且在这种情况下 one-to-one 规则会比 one-to-many 规则的鲁棒性更差。所以我们认为规则应该是 prediction-aware 的。我们首先尝试了 DETR 的思路,直接采用 loss 做 bipartite matching 的 cost[2],发现无论是绝对性能还是有无 NMS 的差距,都得到了进一步的优化。
但我们知道,loss 和 metrics 往往并不一致,它常常要为优化问题做一些妥协(比如做一些加权等等)。也就是说,loss 并不一定是 bipartite matching 的最佳 cost。因而我们提出了一个非常简单的 cost:
看起来稍微有点复杂,但其实就是用网络输出的 prob 代表分类,网络输出和 gt 的 IoU 代表回归,做了加权几何平均,再加一个类似于 inside gt box 的空间先验。加权几何平均和空间先验我们后面都分别做了 ablation。
这就是我们提出的 POTO 策略,它进一步地提升了无 NMS 下的性能,也侧面验证了 loss 并不一定是最好的 cost[3]。但从 Table 1 中我们也发现了,POTO 的性能依旧不能匹敌 one-to-many+NMS 组合。我们认为问题出在两个方面:
1. one-to-one 需要网络输出的 feature 非常 sharp,这对 CNN 提出了较严苛的要求(这也是 Transformer 的优势);
2. one-to-many 带来了更强的监督和更快的收敛速度。
我们分别用 3D Max Filtering 和 one-to-many auxiliary loss 缓解如上问题。
3D Max Filtering
针对第一点,我们提出了 3D Max Filtering,这基于一个 intuition(paper 中没有提到):卷积是线性滤波器,学习 max 操作是比较困难的。此外,我们在 FCOS 做了实验,发现 duplicated prediction 基本来自于 5x5 的邻域内,所以最简单的做法就是在网络中嵌入最常见的非线性滤波器 max pooling。另外,NMS 是所有 feature map 一起做的,但网络在结构上缺少层间的抑制,所以我们希望 max pooling 是跨层的。
如 Figure 3 所示,这个模块只采用了卷积、插值、max pooling 3d,速度非常快,也不需要写 cuda kernel。
One-to-many auxiliary loss
针对第二点监督不够强、收敛速度慢,我们依旧采用 one-to-many assignment 设计了 auxiliary loss 做监督,该 loss 只包含分类 loss,没有回归 loss。assignment 本身没什么可说的,appendix 的实验也表明多种做法都可以 work。这里想提醒大家的是注意看 Figure 2 的乘法,它是 auxiliary loss 可以 work 的关键。在乘法前的一路加上 one-to-many auxiliary loss,乘法后是 one-to-one 的常规 loss。由于 10=0,11=1,我们只需要大致保证 one-to-one assignment 的正样本在 one-to-many 中依然是正样本即可。
实验
最主要的实验结果已经在 Table 1 中呈现了,此外还有一些 ablation 实验。
这里 highlight 几点:
1. α越低,分类权重越大,有无 NMS 的差距越小,但绝对性能也会降低;α太高也不好,我们后续所有实验用α=0.8;
2. 在α合理的情况下,空间先验不是必须的,但空间先验能够在匹配过程中帮助排除不好的区域,提升绝对性能;我们在 COCO 实验中采用 center sampling radius=1.5,在 CrowdHuman 实验中采用 inside gt box;
3. 加权几何平均数(Mul)比加权算术平均数(Add)更好。
去掉 NMS 的最大收益其实是 crowd 场景,这在 COCO 上并不能很好地体现出来。所以我们又在 CrowdHuman 上做了实验如下:
请注意 CrowdHuman 的 ground-truth 做 NMS threshold=0.6,只有 95.1%的 Recall,这也是 NMS 方法的理论上限。而我们的方法没有采用 NMS,于是轻易超越了这一上限。
我们还做了其它一些实验和分析,欢迎看原文。
可视化
经过以上方法,我们成功把 one-to-one 的性能提升到了与 one-to-many+NMS 方法 comparable 的水平。我们可视化了 score map,可以发现 FCN 是有能力学出非常 sharp 的表示的,这也是很让我们惊奇的一点。
结果图中比较明显的改善出现在多峰 case 上。比如两个物体有一定的 overlap(但又没有特别重合),这个时候 one-to-many+NMS 方法经常出现的情况是,除了两个物体分别出了一个框之外,在两个物体中间也出了一个框,这个框与前两个框的 IoU 不足以达到 NMS threshold,但置信度又比较高。这类典型的多峰问题在 POTO 中得到了较大的缓解。
Others
有些人可能比较关心训练时间,因为潜意识里在 dense prediction 上做 bipartite matching 应该是很慢的。然而实际上依赖于 scipy 对linear_sum_assignment的优化,实际训练时间仅仅下降了 10%左右。
如果对这一时间依然敏感,可以用 topk(k=1)代替 bipartite matching;在 dense prediction 里 top1 实际上是 bipartite matching 的近似解。相似地,k>1 的情况对应了 one-to-many 的一种新做法,我们组也对此基于最优传输做了一些工作,后续可能会放出来。
参考
1. 如果有人感兴趣的话,可以在 YOLO 上去掉 NMS 尝试一下,可以接近 30mAP。
2. 注意我们这里没有使用 DETR 的 CE+GIoU+L1 组合,而是直接采用 loss 本身(Focal+GIoU)。我们认为这样更符合 DETR 用 loss 做 cost 的原意。
3. 其实这里可以有一个脑洞留给大家,因为 cost 是不需要求导的,所以甚至是可以直接算 AP 当 cost 的。
4. 侧面印证了分类和回归的冲突在检测任务上是显著的。
5. 事实上加权几何平均数的负对数就是 CE+IoU Loss,加权算术平均数则没有明显的物理含义。
6. NoisyAnchor 在 assign 中采用了类似的公式,只不过采用的是 anchor IoU。
7. 更具体来讲,top1 是 Hugarian Algorithm 只做第一次迭代的结果;由于在 dense prediction 下冲突会很少,一次迭代就已经逼近了最优匹配,这也是为什么 Hungarian Algorithm 这里实际运行很快。
作者介绍 :
王剑锋,北京航空航天大学软件工程硕士
评论