YOLOv5改进系列（14）——更换NMS（非极大抑制）之 DIoU-NMS、CIoU-NMS、EIoU-NMS、GIoU-NMS 、SIoU-NMS、Soft-NMS
☀️一、NMS（非极大抑制）简介
1.1 什么是NMS？
NMS（non maximum suppression）即非极大抑制，顾名思义就是抑制不是极大值的元素，搜索局部的极大值。在最近几年常见的物体检测算法（包括RCNN、SPPNet、Fast-RCNN、Faster-RCNN等）中，最终都会从一张图片中找出很多个可能是物体的矩形框，然后为每个矩形框为做类别分类概率。
如果用一句话概括NMS的意思就是：筛选出一定区域内属于同一种类别得分最大的框。
如下图，网络模型可以给每个检测框一个score，score越大，说明检测框越接近真实值。
然后非最大值抑制的目的是删除score小的框，只剩下sore最大作为最终的预测结果。
1.2 NMS的计算过程
- 定义置信度阈值和IOU阈值取值
- 按置信度降序排列边界框bounding_box
- 从bbox_list中删除置信度小于阈值的预测框
- 循环遍历剩余框，首先挑选置信度最高的框作为候选框
- 接着计算其他和候选框属于同一类的所有预测框和当前候选框的IOU
- 如果上述任两个框的IOU的值大于IOU阈值，那么从box_list中移除置信度较低的预测框
- 重复此操作，直到遍历完列表中的所有预测框
1.3 NMS的局限性
（1）需要手动设置阈值，阈值的设置会直接影响重叠目标的检测，太大造成误检，太小达不到理想情况。
（2）低于阈值的直接设置score为0，做法就比较麻烦了
（3）通过IoU来评估，IoU的做法对目标框尺度和距离的影响不同
1.4 NMS的改进思路
（1）根据手动设置阈值的缺陷，通过自适应的方法在目标稀疏时使用小阈值，目标稠密时使用大阈值。例如Adaptive NMS
（2）由于将低于阈值的直接置为0的做法比较困难，所以我们可以通过将其根据IoU大小来进行惩罚衰减，则变得更加平滑。例如Soft NMS，Softer NMS
（3）IoU的做法存在一定缺陷，改进思路是将目标尺度、距离引进IoU的考虑中。如DIoU等
☀️二、DIoU-NMS、CIoU-NMS、EIoU-NMS、GIoU-NMS 、 SIoU-NMS
🌲2.1 更换DIoU-NMS
一个成熟的IoU衡量指标应该要考虑预测框与真实框的重叠面积、中心点距离、长宽比三个方面。但是IoU 只考虑到了预测框与真实框重叠区域，并没有考虑到中心点距离、长宽比。
DIoU-NMS在DIoUloss一文中提出，不仅仅考虑IoU，还考虑两个框中心点之间的距离。如果两个框之间IoU比较大，但是两个框的中心距离比较大时，可能会认为这是两个物体的框而不会被过滤掉。由于DIoU的计算考虑到了两框中心点位置的信息，故使用DIoU进行评判的nms效果更符合实际，效果更优。
公式：
第①步 修改general.py
在YOLOv5当中，作者是直接调用了Pytorch官方的NMS的API。
也就是在general.py中的non_max_suppression函数中
首先将下面一段代码粘贴到utils/general.py，重新定义NMS模块。这里的计算IoU的函数bbox_iou则是直接引用了YOLOv5中的代码，其简洁的集成了CIoU、SIoU、EIoU、GIoU、DIoU 的计算。
- def NMS(boxes, scores, iou_thres, class_nms='CIoU'):
- # class_nms=class_nms
- GIoU=CIoU=DIoU=EIoU=SIoU=False
- if class_nms == 'CIoU':
- CIoU=True
- elif class_nms == 'DIoU':
- DIoU=True
- elif class_nms == 'GIoU':
- GIoU=True
- elif class_nms == 'EIoU':
- EIoU=True
- else :
- SIoU=True
- B = torch.argsort(scores, dim=-1, descending=True)
- keep = []
- while B.numel() > 0:
- index = B[0]
- keep.append(index)
- if B.numel() == 1: break
- iou = bbox_iou(boxes[index, :], boxes[B[1:], :], GIoU=GIoU, DIoU=DIoU, CIoU=CIoU, EIoU=EIoU, SIoU=SIoU)
- inds = torch.nonzero(iou <= iou_thres).reshape(-1)
- B = B[inds + 1]
- return torch.tensor(keep)
第②步 更换NMS
然后我们将non_max_suppression 函数中的
i = torchvision.ops.nms(boxes, scores, iou_thres)替换为
i = NMS(boxes, scores, iou_thres, class_nms='DIoU')这样就可以还是训练了~
🌲2.2 更换其他的NMS
其余几个方法都是一样的，只需要在第②步改个名称即可：
- DloU：
i = NMS(boxes, scores, iou_thres, class_nms='DIoU') - SloU：
i = NMS(boxes, scores, iou_thres, class_nms='SIoU') - GloU：
i = NMS(boxes, scores, iou_thres, class_nms='GIoU') - EloU：
i = NMS(boxes, scores, iou_thres, class_nms='EIoU') ☀️三、Soft-NMS
根据前面对目标检测中NMS的算法描述，易得出传统NMS的不足：如果一个物体在另一个物体重叠区域出现，即当两个目标框接近时，分数更低的框就会因为与之重叠面积过大而被删掉，从而导致对该物体的检测失败并降低了算法的平均检测率。
上图中，有两匹马，都是待检测目标，也有两个检测到的框，得分分别是0.80和0.95
如果用NMS算法，得分最高的框是红色框，得分0.95，而绿色框与红色框通过计算IoU，肯定是大于一般我们设置的0.5的，那么绿色框就会被删除，导致少检测到一匹马的情况。此外，NMS算法设置阈值也比较麻烦，如果设置过小，那么会出先这样的事情，少检测到目标；如果设置过大，又会经常出先误检。
因此，出现升级版Soft-NMS。具体流程就是我们把NMS算法中去除其他边界框改成，修改其他边界框的置信度。
第①步 修改general.py
同样，首先将下面一段代码粘贴到utils/general.py，重新定义NMS模块。
- def my_soft_nms(bboxes, scores, iou_thresh=0.5, sigma=0.5, score_threshold=0.25):
- bboxes = bboxes.contiguous()
- x1 = bboxes[:, 0]
- y1 = bboxes[:, 1]
- x2 = bboxes[:, 2]
- y2 = bboxes[:, 3]
- # 计算每个box的面积
- areas = (x2 - x1 + 1) * (y2 - y1 + 1)
- # 首先对所有得分进行一次降序排列,仅此一次,以提高后续查找最大值速度. oeder为降序排列后的索引
- _, order = scores.sort(0, descending=True)
- # NMS后,保存留下来的边框
- keep = []
- while order.numel() > 0:
- if order.numel() == 1: # 仅剩最后一个box的索引
- i = order.item()
- keep.append(i)
- break
- else:
- i = order[0].item() # 保留首个得分最大的边框box索引,i为scores中实际坐标
- keep.append(i)
- # 巧妙使用tersor.clamp()函数求取order中当前框[0]之外每一个边框,与当前框[0]的最大值和最小值
- xx1 = x1[order[1:]].clamp(min=x1[i])
- yy1 = y1[order[1:]].clamp(min=y1[i])
- xx2 = x2[order[1:]].clamp(max=x2[i])
- yy2 = y2[order[1:]].clamp(max=y2[i])
- # 求取order中其他每一个边框与当前边框的交集面积
- inter = (xx2 - xx1).clamp(min=0) * (yy2 - yy1).clamp(min=0)
- # 计算order中其他每一个框与当前框的IoU
- iou = inter / (areas[i] + areas[order[1:]] - inter) # 共order.numel()-1个
- idx = (iou > iou_thresh).nonzero().squeeze() # 获取order中IoU大于阈值的其他边框的索引
- if idx.numel() > 0:
- iou = iou[idx]
- newScores = torch.exp(-torch.pow(iou, 2) / sigma) # 计算边框的得分衰减
- scores[order[idx + 1]] *= newScores # 更新那些IoU大于阈值的边框的得分
- newOrder = (scores[order[1:]] > score_threshold).nonzero().squeeze()
- if newOrder.numel() == 0:
- break
- else:
- newScores = scores[order[newOrder + 1]]
- maxScoreIndex = torch.argmax(newScores)
- if maxScoreIndex != 0:
- newOrder[[0, maxScoreIndex],] = newOrder[[maxScoreIndex, 0],]
- # 更新order.
- order = order[newOrder + 1]
- # 返回保留下来的所有边框的索引值,类型torch.LongTensor
- return torch.LongTensor(keep)
第②步 更换NMS
将general.py中将NMS改为soft nms。
这步也是和上面一样，将non_max_suppression 函数中的
i = torchvision.ops.nms(boxes, scores, iou_thres)替换为
i = my_soft_nms(boxes, scores, iou_thres) # 最后就可以开始训练了~