YOLOv11小白的进击之路(六)创新YOLO的iou及损失函数时的源码分析

iou或者说是损失函数的修改经常作为论文的创新点之一,那这篇文章就总结分析了在对YOLO11进行损失函数创新时需要关注的源代码新的一年祝大家论文与财都发发发

总的来看需要关注三个函数,分别位于YOLO庞大源码的不同文件,下面逐一分析:

bbox_iou函数

bbox_iou函数位于/ultralytics-main/ultralytics/utils/metrics.py这个函数的目的是计算两个边界框(box)之间的IoU(Intersection over Union)或其变体(GIoU、DIoU、CIoU)

def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
  • box1 表示单个边界框,box2 表示多个边界框。
  • xywh=True 表示边界框的输入格式是 (x, y, w, h),即中点坐标和宽高,否则将视作 (x1, y1, x2, y2) 格式。
  • GIoUDIoUCIoU 分别对应不同的IoU改进算法。
  • eps 是一个很小的数,用来避免除以零等数值问题。

根据格式确定坐标

    if xywh:  # transform from xywh to xyxy
        (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
        w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
        b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
        b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
    else:  # x1, y1, x2, y2 = box1
        b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
        b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
        w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
        w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
  • 如果 xywh=True,则将中心点 (x, y, w, h) 转为左上角和右下角 (x1, y1, x2, y2)
  • 转换方式:
    • x1 = x - w/2x2 = x + w/2
    • y1 = y - h/2y2 = y + h/2
  • 如果不是 (x, y, w, h) 格式,就直接从 (x1, y1, x2, y2) 中拆分。

计算相交面积 (Intersection)

inter = (  
    (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp_(0) *  
    (b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).clamp_(0)  
)
  • b1_x2.minimum(b2_x2) 相当于取两者右边界的最小值作为相交区域的右边界;b1_x1.maximum(b2_x1) 取两者左边界的最大值作为相交区域的左边界。二者相减得到相交区域的宽度。
  • 同理,对 y 方向做类似操作,计算相交区域的高度。
  • clamp_(0) 会把负值变成 0,代表如果没有重叠(高度或宽度小于 0)则相交面积为 0。

计算联合面积 (Union) 与 IoU

union = w1 * h1 + w2 * h2 - inter + eps  
iou = inter / union
  • 联合面积 = 两个框面积之和 - 相交面积。
  • 计算好的 union 就能用来得到基本的 iou = inter / union
  • + eps 防止分母为 0。

GIoU / DIoU / CIoU(扩展)

如果设置了 GIoU=TrueDIoU=True 或 CIoU=True,则会进行更复杂的计算:

if CIoU or DIoU or GIoU:  
    cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)  # 包覆两个框的最小外接框的宽度  
    ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)  # 同理,高度
  • 将两个框最小能包住它们的“外接矩形”称为“convex box”,用它的宽和高来帮助我们改进IoU的计算。
GIoU
  • c_area = cw * ch 表示两个框的外接矩形面积。
  • GIoU 在 IoU 基础上还会根据这个“外接矩形”对空白区域进行惩罚,从而在没有交集时训练效果更好。
DIoU
if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw.pow(2) + ch.pow(2) + eps  # convex diagonal squared
            rho2 = (
                (b2_x1 + b2_x2 - b1_x1 - b1_x2).pow(2) + (b2_y1 + b2_y2 - b1_y1 - b1_y2).pow(2)
            ) / 4  # center dist**2
            if CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                v = (4 / math.pi**2) * ((w2 / h2).atan() - (w1 / h1).atan()).pow(2)
                with torch.no_grad():
                    alpha = v / (v - iou + (1 + eps))
                return iou - (rho2 / c2 + v * alpha)  # CIoU
            return iou - rho2 / c2  # DIoU
  • DloU还会计算两个框中心点之间的距离,与外接矩形对角线距离比较,用距离来惩罚。
CIoU
  • 在 DIoU 的基础上再加上长宽比等因素,进一步提高回归精度。

最终会返回不同公式下的结果,比如:

  • return iou - (rho2 / c2) 对应 DIoU
  • return iou - (c_area - union) / c_area 对应 GIoU
  • return iou - (rho2 / c2 + v * alpha) 对应 CIoU

简而言之,这个函数的流程是:

  1. 将输入框转换为统一的 (x1, y1, x2, y2) 形式;
  2. 计算相交区域(Intersection);
  3. 计算并返回基本 IoU 或对应改进版本(GIoU、DIoU、CIoU),这些算法通过对距离、外接矩形或长宽比进行额外惩罚,解决 IoU 在一些特殊场景下对回归不敏感的问题。

BboxLoss函数

第二个需要关注的函数是位于/ultralytics-main/ultralytics/utils/loss.py路径下的的BboxLoss函数,主要作用就是先算 IoU 损失做粗定位,再用分布回归进行细定位,二者加权求和后再与其他损失(如分类损失)合并,帮助 YOLO 进行更精准的检测和定位。

类初始化 (init)

class BboxLoss(nn.Module):  
    def __init__(self, reg_max=16):  
        super().__init__()  
        self.dfl_loss = DFLoss(reg_max) if reg_max > 1 else None
  1. reg_max 用于分布回归 (Distribution Focal Loss, DFL) 的参数控制,当 reg_max > 1 时,就会实例化 DFLoss;否则不启用 DFL。
  2. self.dfl_loss 是一个可选的分布回归损失,用于更精细地回归边界框

forward 函数

def forward(self,   
            pred_dist,   
            pred_bboxes,   
            anchor_points,   
            target_bboxes,   
            target_scores,   
            target_scores_sum,   
            fg_mask):  

这个方法接收以下主要输入:

  • pred_dist:分布式预测结果(若启用 DFL,会用来计算分布回归损失)。
  • pred_bboxes:模型预测的边界框坐标,一般是 (x1, y1, x2, y2) 格式。
  • anchor_points:对应特征点或锚框坐标信息,用来辅助计算真值的分布回归标签。
  • target_bboxes:真实的边界框坐标 (Ground Truth)。
  • target_scores:真实样本的分类得分(目标的置信度),或称“目标存在”分数。
  • target_scores_sum:对所有前景框的 target_scores 求和,用于做损失归一化。
  • fg_mask:前景掩码,表示哪些位置是真正需要回归的正样本(Foreground)。

IoU Loss 部分

weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)  
iou = bbox_iou(  
    pred_bboxes[fg_mask],   
    target_bboxes[fg_mask],   
    xywh=False,   
    CIoU=True  
)  
loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum
  1. 计算权重 (weight)

  • target_scores.sum(-1)[fg_mask]:对前景区域的目标分数做叠加,并取前景对应的值。
  • unsqueeze(-1) 是为了在最后一维上加一个新维度,以便与 IoU 逐元素相乘。
  1. 计算 IoU

  • 调用咱们前面分析过的 bbox_iou 函数,这里指定 xywh=False 表示坐标是 (x1, y1, x2, y2)。
  • 使用 CIoU=True,即计算 Complete IoU,对距离和长宽比等额外因素也有一定惩罚。
  1. 组合得到 IoU 损失

  • (1.0 - iou) 作为回归损失的基本形式,IoU 越大,损失越小。
  • * weight 对每个前景实例进行加权。
  • 最后除以 target_scores_sum 做归一化,避免不同批次间正样本数量差异过大。

DFL (Distribution Focal Loss) 部分

if self.dfl_loss:  
    target_ltrb = bbox2dist(anchor_points, target_bboxes, self.dfl_loss.reg_max - 1)  
    loss_dfl = self.dfl_loss(  
        pred_dist[fg_mask].view(-1, self.dfl_loss.reg_max),   
        target_ltrb[fg_mask]  
    ) * weight  
    loss_dfl = loss_dfl.sum() / target_scores_sum  
else:  
    loss_dfl = torch.tensor(0.0).to(pred_dist.device)

1.判断是否启用 DFL

  • 若 self.dfl_loss 不为空,则进入 DFL 分布回归损失的计算;否则直接 loss_dfl = 0.0

2.bbox2dist 函数

  • 就是一个辅助函数,用来把 (anchor_points, target_bboxes) 转换成“分布形式”的真值,用于和 pred_dist(预测的分布)对齐。
  • DFL 会把单一数值(比如 x 的偏移量)表示成一个离散分布,然后用分类方式进行回归,细化了定位精度。

3.计算 DFLoss

  • pred_dist[fg_mask].view(-1, self.dfl_loss.reg_max):把预测的分布展开到合适维度,与真值分布比较。
  • 用 target_ltrb[fg_mask] 做标签,对应四个 (left, top, right, bottom) 或类似的分布表示。
  • 同样乘以 weight 并除以 target_scores_sum,让损失与 IoU Loss 在同一尺度上进行加权和归一化。

最终返回

return loss_iou, loss_dfl
  • 返回 IoU 损失和 DFL 损失,训练时通常会将这两个损失与分类损失等合并,形成完整的损失函数进行反向传播。

整体逻辑是先算 IoU 损失做粗定位,再用分布回归进行细定位,二者加权求和后再与其他损失(如分类损失)合并,帮助 YOLO 进行更精准的检测和定位。

iou_calculation函数

最后咱们需要关注的是位于/ultralytics-main/ultralytics/utils/tal.py路径下的iou_calculation函数,主要作用是进行水平边界框的 IoU计算。

def iou_calculation(self, gt_bboxes, pd_bboxes):  
    """IoU calculation for horizontal bounding boxes."""  
    return bbox_iou(gt_bboxes, pd_bboxes, xywh=False, CIoU=True).squeeze(-1).clamp_(0)

1.调用 bbox_iou 函数

  • 这里把 xywh=False,说明输入的 gt_bboxes 和 pd_bboxes 均是 (x1, y1, x2, y2) 格式。
  • CIoU=True 表示计算的是 Complete IoU,会额外考虑中心点距离及长宽比的惩罚。

2.squeeze(-1)

  • bbox_iou 可能返回形状为 [N, 1] 的张量,这里通过 squeeze(-1) 去除最后一维,使结果变为 [N]
  • clamp_(0)将结果张量中小于 0 的值重置为 0,保证 IoU 的数值不为负数。

最后,得到的就是对于一批水平(轴对齐)边界框的 CIoU 值向量,每个元素对应一对 gt_bbox 和 pd_bbox 的 IoU 结果

至此,我们就完成了对YOLO中iou计算和损失函数的源码分析,YOLOv11小白的进击之路系列持续更新中...

最后

欢迎一起交流探讨 ~ 砥砺奋进,共赴山海!

你可能感兴趣的:(YOLO,YOLO,pytorch,yolo,计算机视觉,人工智能,python)