iou或者说是损失函数的修改经常作为论文的创新点之一,那这篇文章就总结分析了在对YOLO11进行损失函数创新时需要关注的源代码,新的一年祝大家论文与财都发发发!
总的来看需要关注三个函数,分别位于YOLO庞大源码的不同文件,下面逐一分析:
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)
格式。GIoU
、DIoU
、CIoU
分别对应不同的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/2
, x2 = x + w/2
y1 = y - h/2
, y2 = y + h/2
(x, y, w, h)
格式,就直接从 (x1, y1, x2, y2)
中拆分。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 = w1 * h1 + w2 * h2 - inter + eps
iou = inter / union
union
就能用来得到基本的 iou = inter / union
。+ eps
防止分母为 0。如果设置了 GIoU=True
、DIoU=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) # 同理,高度
c_area = cw * ch
表示两个框的外接矩形面积。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
最终会返回不同公式下的结果,比如:
return iou - (rho2 / c2)
对应 DIoU。return iou - (c_area - union) / c_area
对应 GIoU。return iou - (rho2 / c2 + v * alpha)
对应 CIoU。简而言之,这个函数的流程是:
(x1, y1, x2, y2)
形式;第二个需要关注的函数是位于/ultralytics-main/ultralytics/utils/loss.py路径下的的BboxLoss函数,主要作用就是先算 IoU 损失做粗定位,再用分布回归进行细定位,二者加权求和后再与其他损失(如分类损失)合并,帮助 YOLO 进行更精准的检测和定位。
class BboxLoss(nn.Module):
def __init__(self, reg_max=16):
super().__init__()
self.dfl_loss = DFLoss(reg_max) if reg_max > 1 else None
reg_max
用于分布回归 (Distribution Focal Loss, DFL) 的参数控制,当 reg_max > 1
时,就会实例化 DFLoss
;否则不启用 DFL。self.dfl_loss
是一个可选的分布回归损失,用于更精细地回归边界框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)。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
计算权重 (weight)
target_scores.sum(-1)[fg_mask]
:对前景区域的目标分数做叠加,并取前景对应的值。unsqueeze(-1)
是为了在最后一维上加一个新维度,以便与 IoU 逐元素相乘。计算 IoU
bbox_iou
函数,这里指定 xywh=False
表示坐标是 (x1, y1, x2, y2)。CIoU=True
,即计算 Complete IoU,对距离和长宽比等额外因素也有一定惩罚。组合得到 IoU 损失
(1.0 - iou)
作为回归损失的基本形式,IoU 越大,损失越小。* weight
对每个前景实例进行加权。target_scores_sum
做归一化,避免不同批次间正样本数量差异过大。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
(预测的分布)对齐。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 损失做粗定位,再用分布回归进行细定位,二者加权求和后再与其他损失(如分类损失)合并,帮助 YOLO 进行更精准的检测和定位。
最后咱们需要关注的是位于/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小白的进击之路系列持续更新中...
欢迎一起交流探讨 ~ 砥砺奋进,共赴山海!