YOLOv5更改损失函数

#-------------------------------------------------Eiou start-------------------------------------------
# 计算两个框的特定IOU
# def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, EIoU=False, eps=1e-7):
#     # Returns the 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
#
#     iou = inter / union
#     # 目标框IOU损失函数的计算
#     if CIoU or DIoU or GIoU or EIoU:
#         # 两个框的最小闭包区域的width
#         cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)  # convex (smallest enclosing box) width
#         # 两个框的最小闭包区域的height
#         ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex height
#
#         if CIoU or DIoU or EIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
#             # 最小外接矩形 对角线的长度平方
#             c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
#             # 两个框中心点之间距离的平方
#             rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
#                     (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center distance squared
#             if DIoU:
#                 return iou - rho2 / c2  # DIoU
#
#             # CIoU 比DIoU多了限制长宽比的因素：v * alpha
#             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 = v / (v - iou + (1 + eps))
#                 return iou - (rho2 / c2 + v * alpha)
#
#             # EIoU 在CIoU的基础上将纵横比的损失项拆分成预测的宽高分别与最小外接框宽高的差值 加速了收敛提高了回归精度
#             elif EIoU:
#                 rho_w2 = ((b2_x2 - b2_x1) - (b1_x2 - b1_x1)) ** 2
#                 rho_h2 = ((b2_y2 - b2_y1) - (b1_y2 - b1_y1)) ** 2
#                 cw2 = cw ** 2 + eps
#                 ch2 = ch ** 2 + eps
#                 return iou - (rho2 / c2 + rho_w2 / cw2 + rho_h2 / ch2)
#
#         # GIoU https://arxiv.org/pdf/1902.09630.pdf
#         c_area = cw * ch + eps  # convex area
#         return iou - (c_area - union) / c_area
#     return iou  # IoU
#-------------------------------------------------Eiou end-------------------------------------------

第二步：在loss.py里将这段代码改成那段代码

#--------------------------原来的计算iou代码--------------------------------------------------
                iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)  # iou(prediction, target)
                lbox += (1.0 - iou).mean()  # iou loss

改成

iou = bbox_iou(pbox, tbox[i], EIoU=True)                 #EOIOU
if type(iou) is tuple:
   lbox += (iou[1].detach().squeeze() * (1 - iou[0].squeeze())).mean()
   iou = iou[0].squeeze()
else:
   lbox += (1.0 - iou.squeeze()).mean()  # iou loss
   iou = iou.squeeze()

第二步：运行train.py,比较算法更改前后的检测精度。
over！