论文地址:Wise-IoU: Bounding Box Regression Loss with Dynamic Focusing Mechanism

GitHub:https://github.com/Instinct323/wiou

摘要:目标检测作为计算机视觉的核心问题,其检测性能依赖于损失函数的设计。边界框损失函数作为目标检测损失函数的重要组成部分,其良好的定义将为目标检测模型带来显著的性能提升。近年来的研究大多假设训练数据中的示例有较高的质量,致力于强化边界框损失的拟合能力。但我们注意到目标检测训练集中含有低质量示例,如果一味地强化边界框对低质量示例的回归,显然会危害模型检测性能的提升。Focal-EIoU v1 被提出以解决这个问题,但由于其聚焦机制是静态的,并未充分挖掘非单调聚焦机制的潜能。基于这个观点,我们提出了动态非单调的聚焦机制,设计了 Wise-IoU (WIoU)。动态非单调聚焦机制使用“离群度”替代 IoU 对锚框进行质量评估,并提供了明智的梯度增益分配策略。该策略在降低高质量锚框的竞争力的同时,也减小了低质量示例产生的有害梯度。这使得 WIoU 可以聚焦于普通质量的锚框,并提高检测器的整体性能。将WIoU应用于最先进的单级检测器 YOLOv7 时,在 MS-COCO 数据集上的 AP-75 从 53.03% 提升到 54.50%

前言:因为我能使用的算力有限,所以做实验时只在 YOLOv7 上做了。而且因为完整的 MS-COCO 需要更大的参数量,训练一个模型需要 3 天时间,所以我只取了其中四分之一的数据进行训练 (28474 张训练图片)。这篇论文的硬伤在于实验量是严重不足的,所以我也不确保这个方法的普适性。此外,WIoU 是一个依赖超参数的方法,所以建议大家根据自己的实验数据调整超参数,而不是直接使用我的默认参数

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

CIoU、SIoU 的 v2 使用和 WIoU v2 一致的单调聚焦机制,v3 使用和 WIoU v3 一致的动态非单调聚焦机制,详见论文的消融实验

在计算速度上,WIoU 所增加的计算成本主要在于聚焦系数的计算、IoU 损失的均值统计。在实验条件相同时,WIoU 因为没有对纵横比进行计算反而有更快的速度,WIoU 的计算耗时为 CIoU 的 87.2%

在性能提升上,数据集的标注质量越差 (当然差到一定程度就不叫数据集了),WIoU 相对其它边界框损失的表现越好。比如说我在一个检测火焰的比赛里用的 WIoU (那时的初版还比较捞) 使 mAP 提升了 1.70% (相比 CIoU)

这篇博客没有像论文那么详细,大家有什么地方看不明白的欢迎在评论区提问,我会抽时间补上的 ~

现有工作

记锚框为 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失,目标框为 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

IoU 用于度量目标检测任务中预测框与真实框的重叠程度,定义为:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

同时,IoU 有一个致命的缺陷,可以在下面公式中观察到。当边界框之间没有重叠时 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 反向传播的梯度消失。这导致重叠区域的宽度 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 在训练时无法更新

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

现有的工作考虑了许多与包围盒相关的几何因素并构造了惩罚项 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 来解决这个问题,现有的边界框损失都是基于加法的损失,并遵循以下范式:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Distance-IoU

DIoU 将惩罚项定义为中心点连接的归一化长度:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

同时为最小包围框的尺寸 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 提供了负梯度,这将使得 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 增大而阻碍预测框与目标框重叠:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

但不可否认的是,距离度量的确是一个极其有效的解决方案,成为高效边界框损失的必要因子。EIoU 在此基础上加大了对距离度量的惩罚力度,其惩罚项定义为:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Complete-IoU

在 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 的基础上,CIoU 增加了对纵横比一致性的考虑:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

其中的 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 描述了纵横比一致性:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

其中 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 反向传播的梯度满足 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失,也就是 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 不可能为预测框的宽高提供同号的梯度。在前文对 DIoU 的分析中可知 DIoU 会产生负梯度 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失,当这个负梯度与 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 正好抵消时,会导致预测框无法优化。而 CIoU 对纵横比一致性的考虑将打破这种僵局

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Scylla-IoU

Zhora Gevorgyan 证明了中心对齐的边界框会具有更快的收敛速度,以 angle cost、distance cost、shape cost 构造了 SIoU。其中 angle cost 描述了边界框中心连线与 x-y 轴的最小夹角:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

distance cost 描述了两边界框的中心点在x轴和y轴上的归一化距离,其惩罚力度与 angle cost 正相关。distance cost 被定义为:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

shape cost 描述了两边界框的形状差异,当两边界框的尺寸不一致时不为 0。shape cost 被定义为:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 与 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 类似,它们都由 distance cost 和 shape cost 组成:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

本文方法

本文所涉及的聚焦机制有以下几种:

WIoU v1 构造了基于注意力的边界框损失,WIoU v2 和 v3 则是在此基础上通过构造梯度增益 (聚焦系数) 的计算方法来附加聚焦机制

Wise-IoU v1

因为训练数据中难以避免地包含低质量示例,所以如距离、纵横比之类的几何度量都会加剧对低质量示例的惩罚从而使模型的泛化性能下降。好的损失函数应该在锚框与目标框较好地重合时削弱几何度量的惩罚,不过多地干预训练将使模型有更好的泛化能力。在此基础上,我们根据距离度量构建了距离注意力,得到了具有两层注意力机制的 WIoU v1:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

为了防止 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 产生阻碍收敛的梯度,将 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 从计算图 (上标 * 表示此操作) 中分离。因为它有效地消除了阻碍收敛的因素,所以我们没有引入新的度量指标,如纵横比

Wise-IoU v2

Focal Loss 设计了一种针对交叉熵的单调聚焦机制,有效降低了简单示例对损失值的贡献。这使得模型能够聚焦于困难示例,获得分类性能的提升。类似地,我们构造了 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 的单调聚焦系数 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

在模型训练过程中,梯度增益 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 随着 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 的减小而减小,导致训练后期收敛速度较慢。因此,引入 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 的均值作为归一化因子:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

其中的 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 为动量为 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 的滑动平均值,动态更新归一化因子使梯度增益 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 整体保持在较高水平,解决了训练后期收敛速度慢的问题

Wise-IoU v3

定义离群度以描述锚框的质量,其定义为:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

离群度小意味着锚框质量高,我们为其分配一个小的梯度增益,以便使边界框回归聚焦到普通质量的锚框上。对离群度较大的锚框分配较小的梯度增益,将有效防止低质量示例产生较大的有害梯度。我们利用 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 构造了一个非单调聚焦系数并将其应用于 WIoU v1:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

其中,当 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 时,Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 使得 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失。当锚框的离群程度满足 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 (Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 为定值)时,锚框将获得最高的梯度增益。由于 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 是动态的,锚框的质量划分标准也是动态的,这使得 WIoU v3 在每一时刻都能做出最符合当前情况的梯度增益分配策略

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

为了防止低质量锚框在训练初期落后,我们初始化 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 使得 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 的锚框具有最高的梯度增益。为了在训练的早期阶段保持这样的策略,需要设置一个小的动量 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 来延迟 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 接近真实值 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 的时间。对于 batchs 数量为 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 (总数据量 / batch size)、mAP 的提升速度显著放缓的轮次为 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 的训练 (如上图),建议将动量设置为:

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

这种设置使得经过 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失 轮训练后有 Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失。在训练的中后期,WIoU v3 将小梯度增益分配给低质量的锚框以减少有害梯度。同时 WIoU v3 会聚焦于普通质量的锚框,提高模型的定位性能

核心代码

IoU_Cal 类 (在最上面的 Github 链接里) 可以计算现有的边界框损失 (IoU,GIoU,DIoU,CIoU,EIoU,SIoU,WIoU),核心的类变量有:

此外,聚焦机制会对边界框损失的值进行缩放,具体通过实例方法 _scaled_loss 实现

其中,WIoU v3 包含 α, δ 两个超参数,不同的超参数可能适用于不同的模型和数据集,需要自行调整 _scaled_loss 的缺省值以找到最优解

在将 WIoU v3 引进 YOLOv7 时,先在 train_aux.py 中找到损失函数的位置。ComputeLossAuxOTA 是 train 的时候用的,找到其源代码并进行修改

因为 YOLOv7 对模型性能的比较主要利用 utils/metrics 里的 fitness 函数,与损失值无关。而 ComputeLoss 是在 eval 的时候用的,保证不出 bug 就行

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

在初始化函数动一下手脚,指定使用的损失函数

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

再修改 __call__ 函数 (修改的行已用书签标注出) 

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

Wise-IoU 作者导读:基于动态非单调聚焦机制的边界框损失

再找到 bbox_iou 函数的所在位置,修改边界框损失的计算方法

这里因为形参、返回值都和原函数不同,所以要检查 ComputeLoss 中调用这个函数的地方,以防报错

def bbox_iou(box1, box2, type_, x1y1x2y2=True):
    # 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
    # 将边界框信息拼接
    b1 = torch.stack([b1_x1, b1_y1, b1_x2, b1_y2], dim=-1)
    b2 = torch.stack([b2_x1, b2_y1, b2_x2, b2_y2], dim=-1)
    self = IoU_Cal(b1, b2)
    loss = getattr(IoU_Cal, type_)(b1, b2, self=self)
    iou = 1 - self.iou
    return loss, iou

发表回复