Pytorch：目标检测网络-非极大值抑制(NMS)

Pytorch: 目标检测-非极大值抑制(NMS)及其变种

Copyright: Jingmin Wei, Pattern Recognition and Intelligent System, School of Artificial and Intelligence, Huazhong University of Science and Technology

Pytorch专栏教程链接

---

本教程不商用，仅供学习和参考交流使用，如需转载，请联系本人。

Reference

Soft NMS

Softer NMS

IoU-Net

华中科技大学AIA学院-数据科学基础课堂笔记

《深度学习之 Pytorch 物体检测实战》

在前面的章节中，我们对物体检测模型的思想、结构及实现有了一定的了解，但是要想获得较好的检测性能，检测算法的细节处理也极为重要。

在众多的细节处理中，非极大值抑制、样本的不均衡及模型的过拟合这 $3$ 个问题尤为重要，对模型的性能影响极大。这章将主要对非极大值抑制问题进行详细的分析，并给出一些经典的解决方法。

当前的物体检测算法为了保证召回率，对于同一个真实物体往往会有多于 $1$ 个的候选框输出。由于多余的候选框会影响检测精度，因此需要利用 NMS 过滤掉重叠的候选框，得到最佳的预测输出。

NMS 方法简单有效，并且对检测结果至关重要，在物体检测算法中有着广泛的应用。当前有几种较为常见的 NMS 方法，如图所示，首先是最为基本 NMS 方法，利用得分高的边框抑制得分低且重叠程度高的边框。然而基本的 NMS 存在一些缺陷，简单地过滤掉得分低且重叠度高的边框可能会导致漏检等问题。针对此问题陆续产生了一系列改进的方法，如Soft NMS、Softer NMS 及 IoU-Net 等。

基于上述背景，本文将首先介绍 NMS 的基本过程，然后依次讲解 Soft NMS, Softer NMS 及 IoU-Net 的思想与实现方法。

import torch

NMS基本过程

为了保证物体检测的召回率，在 Faster RNN 或者 SSD 网络的计算输出中，通常都会有不止一个候选框对应同一个真实物体。如图左边存在 $3$ 个候选框，但是候选框 A 和 C 对应的是同一个物体，由于 C 的得分比 A 要低，在评测时，C 候选框会被当做一个 False Positive 来看待，从而降低模型精度。实际上由于候选框 A 的质量要比 C 好，理想的输出是 A 而不是 C ，我们希望能够抑制掉候选框 C 。

因此，物体检测网络通常在最后增加一个非极大值抑制操作，即 NMS，将重复冗余的预测去掉，如右图所示。非极大值抑制，顾名思义就是抑制不是极大值的边框，这里的抑制通常是直接去掉冗余的边框。

这个过程涉及以下两个量化指标。

• 预测得分：NMS 假设一个边框的预测得分越高，这个框就要被优先考虑，其他与其重叠超过一定程度的边框要被舍弃，非极大值即是指得分的非极大值。

• IoU：在评价两个边框的重合程度时，NMS 使用了 IoU 这个指标。如果两个边框的 IoU 超过一定阈值时，得分低的边框会被舍弃。阈值通常会取 $0.5$ 或者 $0.7$ 。

NMS 存在一个非常简约的实现方法，算法输入包含了所有预测框的得分、左上点坐标、右下点坐标共 $5$ 个预测量，以及一个设定的 IoU 阈值。具体流程如下:

1. 按照得分，对所有边框进行降序排列，记录下排列的索引 order ，并新建一个列表 keep ，作为最终筛选后的边框索引结果。

2. 将排序后的第一个边框置为当前边框，并将其保留到 keep 中，再求当前边框与剩余所有框的 IoU 。

3. 在 order 中，仅仅保留 IoU 小于设定阈值的索引，重复第 2 步，直到order中仅仅剩余一个边框，则将其保留到 keep 中，退出循环，NMS结束。

def nms(self, bboxes, scores, thresh=0.5):
    # 利用Pytorch实现NMS算法
    x1 = bboxes[:, 0]
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 1]
    y2 = bboxes[:, 1]
    # 计算每个box的面积
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    # 对得分降序排列，order为索引
    _, order = scores.sort(0, descending=True)
    # keep保留了NMS后留下的边框box
    keep = []
    while order.numel() > 0:
        if order.numel() == 1: # 保留框只剩1个
            i = order.item()
            keep.append(i)
            break
        else: # 还有保留框没有NMS
            i = order[0].item() # 保留scores最大的那个框box[i]
            keep.append(i)
        # 利用tensor.clamp函数求取每个框和当前框的最大值和最小值
        # 将输入input张量每个元素的夹紧到区间 [min,max]，并返回结果到一个新张量
        xx1 = x1[order[1: ]].clamp(min=x1[i]) 
        # 左坐标夹紧的最小值为order中scores最大的框的左坐标，对剩余所有order元素进行夹紧操作
        yy1 = y1[order[1: ]].clamp(min=y1[i]) 
        xx2 = x2[order[1: ]].clamp(max=x2[i]) 
        yy2 = y2[order[1: ]].clamp(max=y2[i]) 
        # 求每一个框和当前框重合部分和总共叠加的面积
        inter = (xx2 - xx2).clamp(min=0) * (yy2 - yy1).clamp(min=0)
        union = areas[i] + areas[order[1: ]] - inter
        # 计算每一个框和当前框的IoU
        IoU = inter / union
        # 保留IoU小于threshold的边框索引
        idx = (IoU <= thresh).nonzero().squeeze()
        if idx.numel() == 0:
            break
        # 这里+1是为了补充idx和order之间的索引差
        order = order[idx+1]
    # 返回保留下的所有边框的索引
    return torch.LongTensor(keep)

NMS 方法虽然简单有效，但在更高的物体检测需求下，也存在如下 $4$ 个缺陷:

• 最大的问题就是将得分较低的边框强制性地去掉，如果物体出现较为密集时，本身属于两个物体的边框，其中得分较低的也有可能被抑制掉，从而降低了模型的召回率。

• 阈值难以确定。过高的阈值容易出现大量误检，而过低的阈值则容易降低模型的召回率，这个超参很难确定。

• 将得分作为衡量指标。NMS 简单地将得分作为一个边框的置信度，但在一些情况间下，得分高的边框不一定位置更准，因此这个衡量指标也有待考量。

• 速度：NMS 的实现存在较多的循环步骤，GPU 的并行化实现不是特别容易，尤其是预测框较多时，耗时较多。

抑制得分：Soft NMS

Improving Object Detection With One Line of Code

NMS 方法虽有效过滤了重复框，但也容易将本属于两个物体框中得分低的框抑制掉，从而降低了召回率。以下图为例，假设模型对于当前图像给出了两个预测框，类别都为杯子，分类得分分别为 $0.99$ 与 $0.94$ ，两个框的 IoU 为 $0.6$ 。

按照NMS的方法，由于前面的杯子的分类得分高，并且 IoU 超过了阈值(假设为$0.5$ )，则后面的预测框会被抑制掉。但本身这两个框对应两个不同的真实杯子，因此这种NMS方法会导致漏检的现象。

造成这种现象的原因在于 NMS 的计算公式，如式所示。

$$s_i=\{\begin{array}{l}{s}_i,iou(M,b_i)<N_t\\ 0,iou(M,b_i)⩾N_t\end{array}$$

公式中 $s_i$ 代表了每个边框的得分，$M$ 为当前得分最高的框，$b_i$ 为剩余框的某一个，$N_t$ 为设定的阙值，可以看到当 IoU 大于 $N_t$ 时，该边框的得分直接置 $0$ ，相当于被舍弃掉了，从而有可能造成边框的漏检。

基于此原因，诞生了 Soft NMS 方法，利用一行代码即改进了强硬的 NMS 方法。简而言之，Soft NMS对于 IoU 大于阙值的边框，没有将其得分直接置 $0$ ，而是降低该边框的得分，具体方法如下式所示。

$$s_i=\{\begin{array}{l}{s}_i,iou(M,b_i)<N_t\\ s_i(1-iou(M,b_i)),iou(M,b_i)⩾N_t\end{array}$$

从公式中可以看出，利用边框的得分与 IoU 来确定新的边框得分，如果当前边框与边框 $M$ 的 IoU 超过设定阈值 $N_t$ 时，边框的得分呈线性的衰减。

但是，上式并不是一个连续的函数，当一个边框与 $M$ 的重叠 IoU 超过阈值 $N_t$ 时，其得分会发生跳变，这种跳变会对检测结果产生较大的波动，因此还需要寻找一个更为稳定、连续的得分重置函数，最终 Soft NMS 给出了如下式所示的重置函数。

$$s_i=\{\begin{array}{l}{s}_i{e}^{\frac{iou(M,b_i{)}^2}{\sigma }},\forall b_i\notin D\end{array}$$

采用这种得分衰减的方式，对于某些得分很高的边框来说，在后续的计算中还有可能被作为正确的检测框，而不像 NMS 那样"一棒子打死"，因此可以有效提升模型的召回率。

Soft NMS 的计算复杂度与 NMS 相同，是一种更为通用的非极大值抑制方法，可以将 NMS 看做 Soft NMS 的二值化特例。当然，Soft NMS 也是一种贪心算法，并不能保证找到最优的得分重置映射。经过多种实验证明，Soft NMS 在不影响前向速度的前提下，能够有效提升物体检测的精度。

加权平均：Softer NMS

Bounding Box Regression with Uncertainty for Accurate Object Detection

NMS 与 Soft NMS 算法都使用了预测分类置信度作为衡量指标，即假定分类置信度越高的边框，其位置也更为精准。但很多情况下并非如此，例如下面两种情形：

• 对于一个真实物体，所有的预测边框都不准，那么这时应该选择哪一个？还是综合所有边框得到更为精准的一个结果?
• 具有高分类置信度的边框其位置并不是最精准的。

因此，位置的置信度与分类置信度并不是强相关的关系，直接使用分类置信度作为 NMS 的衡量指标并非是最佳选择。基于此现象，Softer NMS 进一步改进了 NMS 的方法，新增加了一个定位置信度的预测，使得高分类置信度的边框位置变得更加准确，从而有效提升了检测的性能。

首先，为了更加全面地描述边框预测，Softer NMS 方法对预测边框与真实物体做了两个分布假设:

• 真实物体的分布是狄拉克 delta 分布，即标准方差为 $0$ 的高斯分布的极限。
• 预测边框的分布满足高斯分布。

基于这两个假设，Softer NMS 提出了一种基于 KL(Kullback-Leibler) 散度的边框回归损失函数 KL loss 。KL 散度是用来衡量两个概率分布的非对称性衡量，KL 散度越接近于 $0$ ，则两个概率分布越相似。

KL 散度又称为相对熵，信息散度，信息增益，常用在对抗神经网络里。主要用来衡量两个分布的相似度。假设连续随机变量 $x$ ，其概率分布为 $p(x)$ ，模型得到的近似分布为 $q(x)$ 。

公式：对于分布 $p,q$ ，KL 散度为：

$$\begin{array}{rl}KL(p||q)& =-\sum _{i=1}^{n}p(x_i)\log q(x_i)-(-\sum _{i=1}^{n}p(x_i)\log p(x_i))\\ & =\sum _{i=1}^{n}p(x_i)\log \frac{p(x_i)}{q(x_i)}\end{array}$$

具体到边框上，KL Loss 是最小化预测边框的高斯分布与真实物体的狄克拉分布之间的 KL 散度。即预测边框分布越接近于真实物体分布，损失越小。

为了描述边框的预测分布，除了预测位置之外，还需要预测边框的标准差，因此 Softer NMS 提出了如图所示的预测结构。

图中上半部为原始的 Fast RCNN 方法的预测，下半部的网络为 Softer NMS 提出的方法。可以看到，Softer NMS 在原 Fast RCNN 预测的基础上，增加了一个标准差预测分支，从而形成边框的高斯分布，与边框的预测起可以求得 KL 损失，由于公式较为复杂，这里就不再展开描述了。

边框的标准差可以被看做边框的位置置信度，因此 Softer NMS 利用该标准差也改善了 NMS 过程。具体过程大体与 NMS 相同，只不过利用标准差改善了高得分边框的位置坐标，从而使其更为精准。

举个例子，在 NMS 的某次循环中，假设当前边框为 $i$ ，则 Softer NMS 会按照式的方法更新边框 $i$ 的坐标。

$$x{1}_i=\frac{{\sum }_{j}x{1}_j/{\sigma }_{x1,j}^2}{{\sum }_{j}1/{\sigma }_{x1,j}^2}$$

公式中 $j$ 代表与 $i$ 的 IoU 大于设定阈值的边框。可以看出，Softer NMS 对于 loU 大于设定阙值的边框坐标进行了加权平均，希望分类得分高的边框能够利用到周围边框的信息，从而提升其位置的准确度。

总体上，Softer NMS 通过提出的 KL Loss 与加权平均的 NMS 策略，在多个数据集上有效提升了检测边框的位置精度。

定位置信度: loU-Net

Acquisition of Localization Confidence for Accurate Object Detection

在当前的物体检测算法中，物体检测的分类与定位通常被两个分支预测。对于候选框的类别，模型给出了一个类别预测，可以作为分类置信度，然而对于定位而言，回归模块通常只预测了一个边框的转换系数，而缺失了定位的置信度，即框的位置准不准，并没有一个预测结果。

定位置信度的缺失也导致了在前面的 NMS 方法中，只能将分类的预测值作为边框排序的依据，然而在某些场景下，分类预测值高的边框不一定拥有与真实框最接近的位置，因此这种标准不平衡可能会导致更为准确的边框被抑制掉。

基于此，旷视提出了 IoU-Net ，增加了一个预测候选框与真实物体之间的 IoU 分支，并基于此改善了 NMS 过程，进一步提升了检测器的性能。

IoU-Net 的整体结构如图所示，基础架构与原始的 Faster RCNN 类似，使用了 FPN 方法作为基础特征提取模块，然后经过 RoI 的 Pooling 得到固定大小的特征图，利用全连接网络完成最后的多任务预测。

同时，IoU-Net 与 Faster RCNN 也有不同之处，主要有 $3$ 点：

• 在 Head 处增加了一个 IoU 预测的分支，与分类回归分支并行。图中的 Jittered RoIs 模块用于 IoU 分支的训练。
• 基于 IoU 分支的预测值，改善了 NMS 的处理过程。
• 提出了 PrRol-Pooling (Precise RoI Pooling) 方法，进步提升 了感兴趣区域池化的精度。

下面对这 $3$ 个主要的改进点进行详细的介绍。

loU预测分支

IoU 分支用于预测每一个候选框的定位置信度。需要注意的是，在训练时 IoU-Net 通过自动生成候选框的方式来训练 IoU 分支，而不是从 RPN 获取。

具体来讲，Jittered Rols 在训练集的真实物体框上增加随机扰动，生成了一系列候选框，并移除与真实物体框 IoU 小于 $0.5$ 的边框。实验证明这种方法来训练 IoU 分支可以带来更高的性能与稳健性。

IoU 分支也可以方便地集成到当前的物体检测算法中。在整个模型的联合训练时，IoU 预测分支的训练数据需要从每一批的输入图像中单独生成。此外，还需要对 IoU 分支的标签进行归一化，保证其分布在 $[-1,1]$ 区间中。

基于定位置信度的NMS

由于 IoU 预测值可以作为边框定位的置信度，因此可以利用其来改善 NMS 过程。IoU-Net 利用 IoU 的预测值作为边框排列的依据，并抑制掉与当前框 IoU 超过设定阈值的其他候选框。

此外，在 NMS 过程中，IoU-Net 还做了置信度的聚类，即对于匹配到同一真实物体的边框，类别也需要拥有一致的预测值。具体做法是，在 NMS 过程中，当边框 A 抑制边框 B 时，通过下式来更新边框 A 的分类置信度。

$$S_A=\max (S_A,S_B)$$

PrRol-Pooling方法

在 Faster R-CNN 中详细介绍了 RoI Align 的方法，通过采样的方法有效避免了量化操作，减小了 RoI Poling 的误差。但 Align 的方法也存在一个缺点， 即对每一个区域都采取固定数量的采样点，但区域有大有小，都采取同一个数量点，显然不是最优的方法。

以此为出发点，IoU-Net 提出了 PrRoI Pooling 方法，采用积分的方式实现了更为精准的感兴趣区域池化，如图8.6中的右图所示。

与 Rol Align 只采样 $4$ 个点不同，PrRol Pooling 方法将整个区域看做是连续的，采用如图中的积分公式求解每一个区域的池化输出值，区域内的每一个点 $(x,y)$ 都可以通过双线性插值的方法得到。这种方法还有一个好处是其反向传播是连续可导的，因此避免了任何的量化过程。

除了以上 $3$ 点，IoU-Net 还提出了一种优化的方法来解决模型最后边框位置的修正，在此就不展开叙述了。

总体上，IoU-Net 提出了一个 IoU 的预测分支，解决了 NMS 过程中分类置信度与定位置信度之间的不一致，可以与当前的物体检测框架一起端到端地训练，在几乎不影响前向速度的前提下，有效提升了物体检测的精度。