YOLOv5改进之十八：损失函数改进为Alpha-IoU损失函数
前 言：作为当前先进的深度学习目标检测算法YOLOv5，已经集合了大量的trick，但是还是有提高和改进的空间，针对具体应用场景下的检测难点，可以不同的改进方法。此后的系列文章，将重点对YOLOv5的如何改进进行详细的介绍，目的是为了给那些搞科研的同学需要创新点或者搞工程项目的朋友需要达到更好的效果提供自己的微薄帮助和参考。之前改进为EIOU和SIOU，有兴趣的朋友可以看看。
YOLOv5改进之七：损失函数改进_人工智能算法研究院的博客-CSDN博客_yolov5修改损失函数
YOLOv5改进之十：损失函数改进为SIOU_人工智能算法研究院的博客-CSDN博客
解决问题：YOLOv5采用CIOU损失函数，本文采用Alpha-IoU损失函数进行改进，提升检测精度。
原理：
论文：https://arxiv.org/abs/2110.13675
代码： https://github.com/Jacobi93/Alpha-IoU
Bounding box 回归通过预测目标的bbox来定位图像/视频中的目标，这是目标检测、定位和跟踪的基础。例如，最高级的目标检测器通常由一个bbox回归分支和一个分类分支组成，其中bbox回归分支生成用于定位对象进行分类的bbox。在这项工作中，作者探索了更有效的损失函数。早期的目标检测工作使用损失进行bbox回归，而近期的工作直接采用定位性能度量，即Intersection over Union (IoU)作为定位损失。与损失相比，IoU损失对bbox scales是不变的，从而有助于训练更好的检测器。然而，当预测框与Ground truth不重叠时，IoU损失会出现梯度消失问题，导致收敛速度减慢，导致检测器不准确。这激发了几种改进的基于IoU的损失设计，包括Generalized IoU (GIoU)、Distance IoU (DIoU)和Complete IoU (CIoU)。GIoU在IoU损失中引入惩罚项以缓解梯度消失问题，而DIoU和CIoU在惩罚项中考虑了预测框与Ground truth 之间的中心点距离和宽高比。通过在现有的IoU Loss中引入power 变换，提出了一个新的IoU损失函数。首先将Box-Cox变换应用于IoU损失，并将其推广为power IoU loss，记为α。这里进一步简化α为αα，并将其推广到更一般的形式通过加上额外的power正则化项。这使本文所提损失函数能够概括现有的基于IoU的损失，包括GIoU、DIoU和CIoU，到一个新的power IoU损失函数以获得更准确的边界框回归和目标检测。实验结果表明，相对于, α(α>1)增加了high IoU目标的损失和梯度，进而提高了bbox回归精度。当时，它降低了High IoU目标的权重，实验可以看出这会影响BBox的回归精度。power参数α可作为调节α-IoU损失的超参数以满足不同水平的bbox回归精度，其中α >1通过更多地关注High IoU目标来获得高的回归精度(即High IoU阈值)。从经验上表明，α对不同的模型或数据集并不过度敏感，在大多数情况下，α=3表现一贯良好。α-IoU损失家族可以很容易地用于改进检测器的效果，在干净或嘈杂的环境下，不会引入额外的参数，也不增加训练/推理时间。
方 法：
第一步修改general.py，增加SIOU。
def bbox_alpha_iou(box1, box2, x1y1x2y2=False, GIoU=False, DIoU=False, CIoU=False, alpha=2, eps=1e-9):
# Returns tsqrt_he IoU of box1 to box2. box1 is 4, box2 is nx4
box2 = box2.T
# Get the coordinates of bounding boxes
if x1y1x2y2: # x1, y1, x2, y2 = box1
b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
else: # transform from xywh to xyxy
b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
# Intersection area
inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
(torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
# Union Area
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
union = w1 * h1 + w2 * h2 - inter + eps
# change iou into pow(iou+eps)
# iou = inter / union
iou = torch.pow(inter/union + eps, alpha)
beta = 2 * alpha
if GIoU or DIoU or CIoU:
cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
c2 = cw ** beta + ch ** beta + eps # convex diagonal
rho_x = torch.abs(b2_x1 + b2_x2 - b1_x1 - b1_x2)
rho_y = torch.abs(b2_y1 + b2_y2 - b1_y1 - b1_y2)
rho2 = (rho_x ** beta + rho_y ** beta) / (2 ** beta) # center distance
if DIoU:
return iou - rho2 / c2 # DIoU
elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
with torch.no_grad():
alpha_ciou = v / ((1 + eps) - inter / union + v)
# return iou - (rho2 / c2 + v * alpha_ciou) # CIoU
return iou - (rho2 / c2 + torch.pow(v * alpha_ciou + eps, alpha)) # CIoU
else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
# c_area = cw * ch + eps # convex area
# return iou - (c_area - union) / c_area # GIoU
return iou - torch.pow((c_area - union) / c_area + eps, alpha) # GIoU
else:
return iou # torch.log(iou+eps) or iou第二步：将loss.py中边框位置回归损失函数改为Alpha-IoU。
结 果：本人在多个数据集上做了大量实验，针对不同的数据集效果不同，都会有一定的差异，配合请来那个花模型会有提升效果，具体需要自己做实验进行。
预告一下：下一篇内容分享YOLOv5改进之十九。有兴趣的朋友可以关注一下我，有问题可以留言或者私聊我哦
PS：损失函数的改进的方法不仅仅是适用改进YOLOv5，也可以改进其他的YOLO网络，比如YOLOv4、v3等。
最后，希望能互粉一下，做个朋友，一起学习交流。