一文打尽目标检测NMS:效率提升方法总结
加入极市专业CV交流群,与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度 等名校名企视觉开发者互动交流!
同时提供每月大咖直播分享、真实项目需求对接、干货资讯汇总,行业技术交流。关注 极市平台 公众号 ,回复 加群,立刻申请入群~
7月22日,极市平台与重磅邀请到ICML 2020杰出论文一作魏恺轩,为我们深度讲解论文相关工作:免调试即插即用的近端优化算法。请大家锁定直播时间7月22日(周三)20:00。详情戳这里,在极市平台后台回复“62”,即可获取直播链接。本次直播由极市平台和中国图象图形学学会青年工作委员联合组织。
可以看到,NMS由于顺序处理的原因,运算效率较为低下。在笔者的实际项目中,NMS往往能占模型计算总时间的40%甚至更多,极大影响了模型的效率。经过笔者一段时间的调研,关于提升NMS运算速度的方法,在这里也将结合代码进行阶段性总结。
所参考的代码库列举如下:
1. Faster RCNN pytorch (rbg大神) 的 CUDA NMS https://github.com/rbgirshick/py-faster-rcnn
2. YOLACT团队提出的Fast NMS https://github.com/dbolya/yolact
3. CIoU团队提出的Cluster NMS https://github.com/Zzh-tju/CIoU
4. SOLOv2团队提出的Matrix NMS https://github.com/WXinlong/SOLO
5. Torchvision封装的免编译CUDA NMS
加速的关键
在进入正题之前,首先我们应该考察一下NMS运算效率的瓶颈在哪?答案自然是IoU计算及顺序迭代抑制。
假如一张图片中有 个检测框,由于顺序处理的原因,某一个框与其他框计算IoU,最少一次,最多有 次。再加上顺序迭代抑制,NMS算法在计算IoU方面,共需要计算IoU至少 次,最多 次。
因此,要想加速NMS,首当其冲应该要将IoU的计算并行化。这个操作在我们使用IoU-based loss的时候就有,只需计算检测框集合 与自身的IoU即可。检测框集合 事先会按照得分降序排列,也就是说 是最高得分框, 是最低得分框。得到如下这个IoU矩阵:
得益于GPU的并行计算,我们可以一次性得到IoU的全部计算结果。这一步就已经极大地解决了IoU计算繁琐又耗时的问题。代码如下:
def box_iou(boxes1, boxes2):
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
"""
Return intersection-over-union (Jaccard index) of boxes.
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
Arguments:
boxes1 (Tensor[N, 4])
boxes2 (Tensor[M, 4])
Returns:
iou (Tensor[N, M]): the NxM matrix containing the pairwise
IoU values for every element in boxes1 and boxes2
"""
def box_area(box):
# box = 4xn
return (box[2] - box[0]) * (box[3] - box[1])
area1 = box_area(boxes1.t())
area2 = box_area(boxes2.t())
lt = torch.max(boxes1[:, None, :2], boxes2[:, :2]) # [N,M,2]
rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) # [N,M,2]
inter = (rb - lt).clamp(min=0).prod(2) # [N,M]
return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
到这里为止,以上列出的5种NMS都可以做到,从速度上来说CUDA NMS和torchvision NMS相对底层,编译后使用,速度稍快,但必然损失了一些灵活度,后面会讲。(关于CUDA NMS的教程,有兴趣的小伙伴可以参考faster-rcnn源码阅读:nms的CUDA编程,非常详实。
在有了IoU矩阵之后,接下来就是应该要如何利用它来抑制冗余框。
IoU矩阵的妙用
以下列举的三篇文献,可谓是将IoU矩阵玩出了花,从不同的角度发扬光大,在NMS加速方面也确实走在正轨上。(所用的一些符号,笔者进行了统一)
1、Fast NMS
Fast NMS是《YOLACT: Real-time Instance Segmentation》一文的其中一个创新点。由于IoU的对称性,即
接着对
1代表保留,0代表抑制。
代码如下:
def fast_nms(self, boxes, scores, NMS_threshold:float=0.5):
'''
Arguments:
boxes (Tensor[N, 4])
scores (Tensor[N, 1])
Returns:
Fast NMS results
'''
scores, idx = scores.sort(1, descending=True)
boxes = boxes[idx] # 对框按得分降序排列
iou = box_iou(boxes, boxes) # IoU矩阵
iou.triu_(diagonal=1) # 上三角化
keep = iou.max(dim=0)[0] < NMS_threshold # 列最大值向量,二值化
return boxes[keep], scores[keep]
优点:
速度比cython编译加速的Traditional NMS快。(上表截自https://github.com/Zzh-tju/CIoU)
2. 可支持与其他提升精度的NMS方法结合。
缺点:
1. Fast NMS会比Traditional NMS抑制更多的框,性能略微下降。
2. 比CUDA NMS慢,约0.2ms。
这里有必要解释一下,为什么Fast NMS会抑制更多的框?
我们知道NMS的思想是:当一个框是冗余框,被抑制后,将不会对其他框产生任何影响。但在Fast NMS中,如果一个框
也就是在刚刚的例子中,由于第二个框
不过呢,YOLACT主要针对的是实例分割,mask是从box中裁剪出来的,Fast NMS对mask AP的下降比较轻微,约0.1~0.3的AP,但似乎对目标检测的box AP会下降更多。
2、Cluster NMS
Cluster NMS出自《Enhancing Geometric Factors in Model Learning and Inference for Object Detection and Instance Segmentation》一文。研究者主要旨在弥补Fast NMS的性能下降,期望也利用pytorch的GPU矩阵运算进行NMS,但同时又使得性能保持与Traditional NMS相同。
最开始看到这个名字时,笔者还以为是采取聚类的NMS,这不禁让人想起了今年2月的FeatureNMS。但后来仔细看了过后,发现这里的cluster的含义不一样。
以上是论文的原话,意思是说,cluster是一个框集合,若某一个框A属于这个cluster,则必有cluster中的框与A的IoU≥NMS阈值,且A不会与除这个cluster之外的框有IoU≥NMS阈值。举个简单的例子:
上图中的黑红蓝橙四个框构成一个cluster,而绿色的两个框构成一个cluster,虽然两个cluster之间有相交,但都没有超过NMS阈值,于是这两个框集合不能合并成一个cluster。
然后我们看一下Cluster NMS是怎么做的?
其实就是一个迭代式的Fast NMS。前面的过程与Fast NMS一模一样,都是
这里借用一下作者在github的流程图。
以笔者的角度看,这种利用矩阵左乘的方式,其实就是在省略上一次Fast NMS迭代中被抑制的框对其他框的影响。因为我们知道一个对角矩阵左乘一个矩阵,就是在做行变换啊,对应行是1,乘完的结果这一行保持不变,如果对应行是0,乘完的结果就是把这一行全部变成0。而
论文里还给出了对这个命题的证明,又是数学数学数学!大致一看是利用数学归纳法证的,感兴趣的可以去看论文,这里请允许我跳过。
def cluster_nms(self, boxes, scores, NMS_threshold:float=0.5):
'''
Arguments:
boxes (Tensor[N, 4])
scores (Tensor[N, 1])
Returns:
Fast NMS results
'''
scores, idx = scores.sort(1, descending=True)
boxes = boxes[idx] # 对框按得分降序排列
iou = box_iou(boxes, boxes).triu_(diagonal=1) # IoU矩阵,上三角化
C = iou
for i in range(200):
A=C
maxA = A.max(dim=0)[0] # 列最大值向量
E = (maxA < NMS_threshold).float().unsqueeze(1).expand_as(A) # 对角矩阵E的替代
C = iou.mul(E) # 按元素相乘
if A.equal(C)==True: # 终止条件
break
keep = maxA < NMS_threshold # 列最大值向量,二值化
return boxes[keep], scores[keep]
好了,至此Cluster NMS算是完成了对Fast NMS性能下降的弥补。我们直接看结果好了
嗯,效果还是可以的,保持了AP与AR一样,运算效率比Fast NMS下降了一些,毕竟是迭代Fast NMS的操作,但也比Traditional NMS快多了。
以为这样就完了?接下来才是重头戏!
这里又分为两个部分,一个是实用上的,另一个是理论上的。先说一下实用上的。
之前说过基于pytorch的NMS方法灵活度要比CUDA NMS更高,就在于这些基于pytorch的NMS是高层语言编写,专为研究人员开发,矩阵运算清晰简洁,于是乎可以很方便地与一些能够提升精度的NMS方法结合!正所谓强强联合,于是就诞生出了又快又好的一系列Cluster NMS的变体。
论文里主要举了三种变体:
得分惩罚机制SPM(Score Penalty Mechanism)
也就是对刚刚Cluster NMS终止后的矩阵
def SPM_cluster_nms(self, boxes, scores, NMS_threshold:float=0.5):
'''
Arguments:
boxes (Tensor[N, 4])
scores (Tensor[N, 1])
Returns:
Fast NMS results
'''
scores, idx = scores.sort(1, descending=True)
boxes = boxes[idx] # 对框按得分降序排列
iou = box_iou(boxes, boxes).triu_(diagonal=1) # IoU矩阵,上三角化
C = iou
for i in range(200):
A=C
maxA = A.max(dim=0)[0] # 列最大值向量
E = (maxA < NMS_threshold).float().unsqueeze(1).expand_as(A) # 对角矩阵E的替代
C = iou.mul(E) # 按元素相乘
if A.equal(C)==True: # 终止条件
break
scores = torch.prod(torch.exp(-C**2/0.2),0)*scores #惩罚得分
keep = scores > 0.01 #得分阈值筛选
return boxes[keep], scores[keep]
2. +中心点距离
说到底DIoU就是该团队提出来的,那怎能不用上中心点距离呢?于是论文在两个方面添加了中心点距离。一个是IoU矩阵直接变成DIoU矩阵,由于DIoU也是满足尺度不变性的,所以完全没问题,相距很远的框之间的DIoU会变成负数,不过不影响过程。第二个是基于上面的SPM,公式如下:
这里多添了一个
加了这两个变体之后,AP与AR得到了明显的改善,速度也是略微下降,很香
3. 加权平均法Weighted NMS
也是利用矩阵
就可以更新框的坐标了。
在YOLOv3上的效果有了不错的改进,虽然速度不及torchvision NMS,增加了5ms的运算成本,但结合Weighted NMS与DIoU,可以提升精度 (最后一行)。
def Weighted_cluster_nms(self, boxes, scores, NMS_threshold:float=0.5):
'''
Arguments:
boxes (Tensor[N, 4])
scores (Tensor[N, 1])
Returns:
Fast NMS results
'''
scores, idx = scores.sort(1, descending=True)
boxes = boxes[idx] # 对框按得分降序排列
iou = box_iou(boxes, boxes).triu_(diagonal=1) # IoU矩阵,上三角化
C = iou
for i in range(200):
A=C
maxA = A.max(dim=0)[0] # 列最大值向量
E = (maxA < NMS_threshold).float().unsqueeze(1).expand_as(A) # 对角矩阵E的替代
C = iou.mul(E) # 按元素相乘
if A.equal(C)==True: # 终止条件
break
keep = maxA < NMS_threshold # 列最大值向量,二值化
weights = (C*(C>NMS_threshold).float() + torch.eye(n).cuda()) * (scores.reshape((1,n)))
xx1 = boxes[:,0].expand(n,n)
yy1 = boxes[:,1].expand(n,n)
xx2 = boxes[:,2].expand(n,n)
yy2 = boxes[:,3].expand(n,n)
weightsum=weights.sum(dim=1) # 坐标加权平均
xx1 = (xx1*weights).sum(dim=1)/(weightsum)
yy1 = (yy1*weights).sum(dim=1)/(weightsum)
xx2 = (xx2*weights).sum(dim=1)/(weightsum)
yy2 = (yy2*weights).sum(dim=1)/(weightsum)
boxes = torch.stack([xx1, yy1, xx2, yy2], 1)
return boxes[keep], scores[keep]
接下来说一下Cluster NMS理论上的好处。
Cluster NMS的迭代次数通常少于Traditional NMS的迭代次数。
这一优点,从理论上给了它使用CUDA编程更进一步加速的可能。大致意思是说,平常我们在做NMS时,迭代都是顺序处理每一个cluster的。在Traditional NMS中,虽然不同的cluster之间本应毫无关系,但计算IoU重复计算了属于不同cluster之间的框,顺序迭代抑制的迭代次数也仍然保持不变。
但在Cluster NMS中,使用这种行变换的方式,就可以将本应在所有cluster上迭代,化简为只需在一个拥有框数量最多的cluster上迭代就够了(妙啊)。不同的cluster间享有相同的矩阵操作,且它们互不影响。这导致迭代次数最多不超过一张图中最大cluster所拥有的框的个数。也就是下面这种情形
上图中共有10个检测框,分成了3个cluster,它的IoU矩阵(被NMS阈值二值化了)在右边。Cluster NMS做迭代的次数最多不超过4次,因为上图中框数量最多的那个cluster(红色)一共只有4个框。而实际上这张图使用Cluster NMS,只需迭代2轮便结束了。
因此这带来的好处是时间复杂度的下降。特别是对于一张图中有很多个cluster时,效果更为显著。比如这种情形,密密麻麻的狗狗
物体数量越多,到时候的检测框也就越多,形成的cluster必然也会增多,于是乎Cluster NMS这种对所有cluster并行处理的算法必然迭代次数非常少,不会随着物体的增多而过分地增加迭代轮数。最极端的情形是
优点:
Cluster NMS可以保持与Traditional NMS一样的精度,速度快于Cython编译的Traditional NMS。 可以结合一些提升精度的方法,比如得分惩罚法,加权平均法,+中心点距离法。 并行处理cluster,给了算法迭代次数一个上界 (最大cluster的框数量)。
缺点:
速度上,各种变体Cluster NMS < Cluster NMS < Fast NMS < CUDA NMS3、Matrix NMSMatrix NMS出自《SOLOv2: Dynamic, Faster and Stronger》,也是利用排序过后的上三角化的IoU矩阵。不同在于它针对的是mask IoU。我们知道mask IoU的计算会比box IoU计算更耗时一些,特别是在Traditional NMS中,会加剧运算开销。因此Matrix NMS将mask IoU并行化是它最大的一个贡献。然后论文同样采取了得分惩罚机制来抑制冗余mask。具体代码如下:
不过笔者在实现这个代码时遇到了问题,惩罚因子decay通常会≥1,在考虑一个mask
随后decay取列最小值。但依据如上代码,decay通常会≥1,这会增大mask的得分,而代码又以0.05作为score筛选的阈值,说明应该是减少得分的思想没有错,这个矛盾导致得出的结果奇差无比。如有任何人知道Matrix NMS出现这种问题的原因,还请告知我问题在哪。
这里直接贴出论文结果,从这里似乎也看出了相比于Soft-NMS每次迭代重排序的做法,直接得分惩罚也是可取的。
优点:
实现了mask IoU的并行计算,对于box-free的密集预测实例分割模型很有使用价值。 与Fast NMS一样,只需一次迭代。
缺点:
与Fast NMS一样,直接从上三角IoU矩阵出发,可能造成过多抑制。总结
1. CUDA NMS与Torchvision NMS稍快于以上三种基于pytorch实现的NMS,但比较死板,改动不易。
2. Cluster NMS基本可以取代Fast NMS,且与其他提升精度的方法的有机结合,使它又快又好。(目前作者团队的代码库只提供了得分惩罚法,+中心点距离,加权平均法三种。)
3. Matrix NMS也可以参考Cluster NMS的方式先得到一个与Traditional NMS一样的结果后,再进行后续处理。
4. 三种基于pytorch的NMS方法,只依赖于矩阵操作,无需编译。
5. 将2D检测的NMS加速方法推广至3D检测,应该也是很有价值的。
参考资料
YOLACT: Real-time Instance Segmentation. ICCV 2019 Enhancing Geometric Factors in Model Learning and Inference for Object Detection and Instance Segmentation. Arxiv 2020.05 SOLOv2: Dynamic, Faster and Stronger. Arxiv 2020.03
推荐阅读