OneNet：一阶段端到端目标检测网络，无需NMS！无需二分匹配！

本文作者：孙培泽 | 编辑：Amusi
https://zhuanlan.zhihu.com/p/331590601
本文已由原作者授权，不得擅自二次转载

本文主要介绍一下我们最近的工作：

OneNet: End-to-End One-Stage Object Detection by Classification Cost

• 论文：https://peizesun.github.io/OneNet.pdf
• 项目代码：https://github.com/PeizeSun/OneNet

我们提出了OneNet，首次实现了在dense detector中无需Non-Maximum Suppression(NMS)后处理。OneNet的样本匹配策略是one-to-one，即一个gt一个正样本，其他都是负样本；正样本是所有样本中和gt的cost最小的样本，cost定义为样本与gt的分类cost(loss)和定位cost(loss)之和。我们发现，分类cost是去除NMS的关键；而回顾之前的dense detector样本匹配策略，都是只有定位cost，导致冗余的高分检测框，如box IoU(e.g.,YOLO, RetinaNet), point distance(e.g.,FCOS, CenterNet)。OneNet在标准的COCO benchmark上使用ResNet-50 单模型达到了35.0 AP和 67 FPS。

1. 简介

现有的end-to-end的目标检测模型都是two-stage或者multiple-stage。这些模型的检测性能很好，但是one-stage在工业应用中有着更大的潜力。我们提出了OneNet：end-to-end one-stage object detector。OneNet的优势是：

• 整个网络是全卷积的。
• 无需Non-Maximum Suppression(NMS)后处理或者self-attention模块。
• 样本匹配策略是简单的Minimum Cost，无需启发式规则或者复杂的最优二分匹配。

cost定义为样本与gt的classification cost(loss)和location cost(loss)之和。我们发现，classification cost是实现end-to-end的关键。而回顾之前的dense detector样本匹配策略，都是只有定位cost，如box IoU(e.g.,YOLO, RetinaNet), point distance(e.g.,FCOS, CenterNet)。只有定位cost的样本匹配策略会导致冗余的高分检测框(图2)，需要NMS后处理去除这些冗余框。

2. OneNet

输入图片(H×W×3)，backbone产生feature map(H/4×W/4×C)，head预测分类(H/4×W/4×K)和回归(H/4×W/4×4), 最后的输出直接取top-k高分框。

图1 OneNet pipeline

Backbone: Backbone是先bottom-up再top-down的结构。其中，bottom-up结构是resnet，top-down结构是FPN。我们实现了两种FPN，一种是上采样中引入deformable conv，为了追求较高检测精度；一种是普通conv，方便工业部署。

Head: Head是两个并行的conv，分类conv预测类别，回归conv预测到物体框的4个边界的距离。

Output: 直接取top-k高分框，没有NMS，也没有类似CenterNet中max-pooling的操作。

样本匹配策略：OneNet的样本匹配策略是一种基于minimum cost的异常简单的方法，没有启发式规则，也没有最优二分匹配。cost定义为样本与gt的classification cost(loss)和location cost(loss)之和，具体定义是：

L_cls是分类focal loss， L_l1和L_giou是预测框和gt框归一化后的l1 loss和giou loss。lambda是系数。

对每个gt，正样本是和gt的cost最小的样本，其他都是负样本。伪代码：

#C is cost matrix, shape of (nr_sample, nr_gt)
C = cost_class + cost_l1 + cost_giou

#Minimum cost, src_ind is index of positive sample
_, src_ind = torch.min(C, dim=0)
tgt_ind = torch.arange(nr_gt)

indices.append((src_ind, tgt_ind))

3. 实验

表1 Ablation studies on matching cost and NMS.

表1中的4个实验都是one-to-one的样本匹配策略。其中第一个实验的location cost是feature map中point的位置到物体gt center的位置的距离(可以理解为CenterNet只有高斯极值点为1，其他都是0)。

从表1可以看出，classification cost是去掉NMS的关键。而回想绝大多数的样本匹配策略，如，box IoU，point distance，都是只考虑了location cost。第三个实验如此拉胯的原因可能是因为predicted box是变化的，会导致正负样本来回横跳，训练低效。

图2给出了4个实验的可视化图，可以看到，没有classification cost的模型会预测出冗余的高分检测框，需要NMS后处理来去除这些冗余框。而引入classification cost的模型消除了冗余框。

图2 表1中的4个对比实验的可视化图

4. 讨论

一年前，在anchor-free和label assignment的那波研究中，曾经考虑过one-to-one的样本匹配策略，表1第一行的实验也做过了，性能也是只有AP 20+(加上NMS 30+)。当时自我解释的原因是：某一位置如果分类是1，那么这一位置的周围位置很难突变成0。e2e one-stage也一直没有做work。

最近DETR出现了，在样本匹配中同时考虑location cost和classification cost，成功做出了e2e two-stage(multiple-stage)。这给人启发e2e one-stage是不是也需要引入classification cost。表1第一行的实验简单地加上classification cost（即表1第二行实验）竟然神奇地work了！甚至optimal bipartite matching也不需要，直接全图找最小cost的样本就行。可能optimal bipartite matching也可以做，但是在dense detector中太慢了。

我们的实验结论提出了很多较为本质的问题：为什么引入classification cost能够使得相邻的feature map points的分类发生突变？样本之间的交互（例如self-attetion）对于e2e是否必须？这些问题都值得后续深入研究。

OneNet 论文和项目代码下载

在CVer微信公众号后台回复：OneNet，即可下载上述论文PDF和项目源代码

重磅！CVer-目标检测 微信交流群已成立

扫码添加CVer助手，可申请加入CVer-目标检测 微信交流群，目前已汇集4000人！涵盖2D/3D目标检测、小目标检测、遥感目标检测等。互相交流，一起进步！