yolo v5 损失函数分析

yolo v5 损失函数分析

    与 yolo v1 类似，v5 损失函数由 3 个部分组成，分别为 bbox 回归损失、目标置信度损失和类别损失。网络每个将特征图分为若干个 cell，每个 cell 输出一个 $[t_x,t_y,t_w,t_h,p_o,c_1,c_2,...]$ 的向量，其中 $t_x,t_y$ 用于计算预测框和对应 anchor box (也就是所在 cell) 两者中心的偏移量，$t_w,t_h$ 用于计算预测框的宽高，$p_o$ 是该 cell (预测框) 含有目标的概率，$c_1,c_2,...$ 为对应类别的预测值。

    三个部分的损失均是通过匹配到的正样本对来计算，每一个输出特征图相互独立，直接相加得到最终每一部分的损失值。先给出整体的计算公式：
$${\mathcal{L}}_{\text{v}5}\left({\mathbf{t}}_{\text{p}},{\mathbf{t}}_{\text{gt}}\right)=\sum _{k=0}^K\left[{\alpha }_k^{\text{balance}}{\alpha }_{\text{box}}\sum _{i=0}^{S^2}\sum _{j=0}^B{\mathbb{I}}_{kij}^{\text{obj}}{\mathcal{L}}_{\text{CIoU}}+{\alpha }_{\text{obj}}\sum _{i=0}^{S^2}\sum _{j=0}^B{\mathbb{I}}_{kij}^{\text{obj}}{\mathcal{L}}_{\text{obj}}+{\alpha }_{\text{cls}}\sum _{i=0}^{S^2}\sum _{j=0}^B{\mathbb{I}}_{kij}^{\text{obj}}{\mathcal{L}}_{\text{cls}}\right]$$
其中，$K,S^2,B$ 分别为输出特征图、cell 和 每个 cell 上 anchor 的数量；${\alpha }_{\star }$ 为对应项的权重，在 hyp.scratch-high.yaml 中默认取值为 ${\alpha }_{\text{box}}=0.05，{\alpha }_{\text{cls}}=0.3,{\alpha }_{\text{obj}}=0.7$; ${\mathbb{I}}_{kij}^{\text{obj}}$ 表示第 $k$ 个输出特征图，第 $i$ 个 cell, 第 $j$ 个 anchor box 是否是正样本，如果是正样本则为 1，反之为 0；${\mathbf{t}}_{\text{p}},{\mathbf{t}}_{\text{p}}$ 是预测向量和 ground-truth 向量；${\alpha }_k^{\text{balance}}$ 用于平衡每个尺度的输出特征图的权重，默认取值为 $[4.0,1.0,0.4]$, 依次对应 $80\times 80,40\times 40,20\times 20$ 的输出特征图。

1. bbox 回归损失

    v5 使用的是 CIoU Loss。
    yolo v5 中正样本匹配策略和 bbox 回归如下图所示。

具体 CIoU Loss 分析可以参考 基于IOU的损失函数合集。

iou_term = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)
lbox += (1.0 - iou_term).mean()

2. 目标置信度损失

    目标置信度损失由正样本匹配得到的样本对计算，一是预测框中的目标置信度分数 $p_o$；二是预测框和与之对应的目标框的 iou 值，其作为 ground-truth。两者计算二进制交叉熵得到最终的目标置信度损失。公式如下：
$${\mathcal{L}}_{\text{obj}}\left(p_o,p_{\text{iou}}\right)={\text{BCE}}_{\text{obj}}^{\text{sig}}\left(p_o,p_{\text{iou}};w_{\text{obj}}\right)$$

BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
obji = self.BCEobj(pi[..., 4], tobj)

3. 类别损失

    类别损失与置信度损失类似，通过预测框的类别分数和目标框类别的 one-hot 表现来计算类别损失，公式如下：
$${\mathcal{L}}_{\text{cls}}\left({\mathbf{c}}_{\text{p}},{\mathbf{c}}_{\text{gt}}\right)={\text{BCE}}_{\text{cls}}^{\text{sig}}\left({\mathbf{c}}_{\text{p}},{\mathbf{c}}_{\text{gt}};w_{\text{cls}}\right)$$

    这里目标置信度损失和类别损失使用的是带 sigmoid 的二进制交叉熵函数 BCEWithLogitsLoss。如果要使用 Focal Loss 在其基础上改动即可。

BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
lcls += self.BCEcls(pi[..., 5:], t_cls)

    源程序分析下次再说。