【NMS】|提高效率

最新推荐文章于 2024-08-07 14:17:28 发布

rrr2

最新推荐文章于 2024-08-07 14:17:28 发布

阅读量506

点赞数

分类专栏：目标检测

本文链接：https://blog.csdn.net/qq_35608277/article/details/118437805

版权

目标检测专栏收录该内容

90 篇文章 5 订阅

订阅专栏

https://zhuanlan.zhihu.com/p/157900024

NMS由于顺序处理的原因，运算效率较为低下。在笔者的实际项目中，NMS往往能占模型计算总时间的40%甚至更多，极大影响了模型的效率。经过笔者一段时间的调研，关于提升NMS运算速度的方法，在这里也将结合代码进行阶段性总结。

所参考的代码库列举如下：

Faster RCNN pytorch (rbg大神) 的 CUDA NMS https://github.com/rbgirshick/py-faster-rcnn
YOLACT团队提出的Fast NMS https://github.com/dbolya/yolact
CIoU团队提出的Cluster NMS https://github.com/Zzh-tju/CIoU
SOLOv2团队提出的Matrix NMS https://github.com/WXinlong/SOLO
Torchvision封装的免编译CUDA NMS

1 fast nms

假如一张图片中有 n个检测框，。检测框集合事先会按照得分降序排列，得到如下这个IoU矩阵X：
在这里插入图片描述

Fast NMS首先对X使用pytorch的triu函数进行上三角化，得到了一个对角线元素及下三角元素都为0的IoU矩阵X。

在这里插入图片描述
接着X执行按列取最大值操作，得到一维张量b，bi表示第i个框与前面i-1个框重复的最大iou值。

最后使用NMS阈值，论文取0.5，对 b 二值化。

元素小于NMS阈值的是保留的框，≥NMS阈值的是冗余框。例如

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]