File size: 3,005 Bytes
6a5bb7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """Wise-IoU v1 动态聚焦权重 — 正确的 ultralytics 8.4.46 集成方式
ultralytics 8.4.46 的 BboxLoss.forward 直接用 iou 标量：
loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum
不支持元组返回。因此 patch BboxLoss.forward，在 weight 中融入 β。
"""
import torch
import torch.nn.functional as F
def _wise_beta(box1, box2, eps=1e-7):
"""计算宽高比偏差动态聚焦系数 β，归一化使 batch 均值≈1。"""
b1_x1, b1_y1, b1_x2, b1_y2 = box1.unbind(-1)
b2_x1, b2_y1, b2_x2, b2_y2 = box2.unbind(-1)
w1 = (b1_x2 - b1_x1).clamp(min=eps)
h1 = (b1_y2 - b1_y1).clamp(min=eps)
w2 = (b2_x2 - b2_x1).clamp(min=eps)
h2 = (b2_y2 - b2_y1).clamp(min=eps)
rw = (w1 / w2).clamp(0.1, 10)
rh = (h1 / h2).clamp(0.1, 10)
raw = (rw - 1) ** 2 + (rh - 1) ** 2
return (raw / (raw.mean() + eps)).detach().unsqueeze(-1) # (N, 1)
def patch_wise_iou():
"""Patch BboxLoss.forward 以融入 Wise-IoU β 动态权重。
只在 BboxLoss.forward 中把 weight 乘以 β，其余逻辑不变。
不替换 bbox_iou，依然使用 CIoU 作为 iou 值（更稳定）。
"""
import ultralytics.utils.loss as los
from ultralytics.utils.loss import bbox_iou, bbox2dist, DFLoss
from ultralytics.utils.tal import dist2bbox
def _wise_forward(
self,
pred_dist,
pred_bboxes,
anchor_points,
target_bboxes,
target_scores,
target_scores_sum,
fg_mask,
imgsz,
stride,
):
weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
beta = _wise_beta(pred_bboxes[fg_mask], target_bboxes[fg_mask])
weighted = weight * beta # Wise-IoU 动态权重
iou = bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True)
loss_iou = ((1.0 - iou) * weighted).sum() / 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:
target_ltrb = bbox2dist(anchor_points, target_bboxes)
target_ltrb = target_ltrb * stride
target_ltrb[..., 0::2] /= imgsz[1]
target_ltrb[..., 1::2] /= imgsz[0]
pred_dist = pred_dist * stride
pred_dist[..., 0::2] /= imgsz[1]
pred_dist[..., 1::2] /= imgsz[0]
loss_dfl = (
F.l1_loss(pred_dist[fg_mask], target_ltrb[fg_mask], reduction="none").mean(-1, keepdim=True) * weight
)
loss_dfl = loss_dfl.sum() / target_scores_sum
return loss_iou, loss_dfl
los.BboxLoss.forward = _wise_forward
print("[WiseIoU] BboxLoss.forward 已替换为 Wise-IoU v1 动态权重版本")
|