yolov5损失函数详解【附代码】

本文章将结合代码对yolov5损失函数部分进行详细说明，包含其中的样本匹配问题。如果还需要学习关于yolov5其他部分内容，可以参考我其他文章。

yolov5语义分割：

            yolov5图像分割中的NMS处理

            yolov5图像分割Segmentation函数

yolov5 trt 学习：

            yolov5 trt推理【python版】

yolov5剪枝：

            yolov5剪枝

通过yaml修改yolov5：

            通过yaml修改yolov5网络

iou样本匹配：

            以yolov4 loss为例讲解iou样本匹配

            NMS预测框的筛选

下面的3个feature_map是仿照v5的head随机产生的输出。为了方便后面代码讲解，这里我设置的num_classes为1 。

feature_map1 = torch.rand([batch_size, 3, 80, 80, 5 + num_classes])
feature_map2 = torch.rand([batch_size, 3, 40, 40, 5 + num_classes])
feature_map3 = torch.rand([batch_size, 3, 20, 20, 5 + num_classes])
pred = [feature_map1, feature_map2, feature_map3]

而标签target如下【也就是通过dataloader处理后的数据集】

我这里的target也是我随便举例的，可以看到他的shape为[3,6]，也就是【num_obj, batch_idex+classes+xywh】。这里的num_obj表示当前图像中有出现了几个目标，batch_idex是第几个图像或者说第几个batch的索引，比如我这里batch是2，但这个是第一张图像的target信息，class表示当前目标是什么类【注意和num_classes区分】，后面的xywh就是box信息。

targets = torch.tensor([[0.00000, 0.00000, 0.04204, 0.21125, 0.08408, 0.36503],
                        [0.00000, 0.00000, 0.14960, 0.24400, 0.23867, 0.36503],
                        [0.00000, 0.00000, 0.36253, 0.24517, 0.21995, 0.39545]])

ComputeLoss

下面这一部分是computeLoss定义的初始化参数。

yolo系列的损失函数通常为三个部分。cls_loss[分类]，obj_loss[置信度loss]，loc_loss[box loss].

前两者从代码中可以看到采用二分类交叉熵。【这里不说focalLoss】

self.balance可以理解为三个head部分的权重，即分配给小中大目标的权重[8080head预测小目标，4040预测中目标，20*20预测小目标]。

na:表示anchors的数量

nc:num_classes

nl:head的数量

class ComputeLoss:
    sort_obj_iou = False

    # Compute losses
    def __init__(self, model, autobalance=False):
        device = next(model.parameters()).device  # get model device
        h = model.hyp  # 获取超参数

        # 损失函数定义，cls:二分类交叉熵， obj:二分类交叉熵
        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))

        # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
        self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0))  # positive, negative BCE targets

        # Focal loss
        g = h['fl_gamma']  # focal loss gamma
        if g > 0:
            BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)

        m = de_parallel(model).model[-1]  # Detect() module
        # balance用于判断是否输出为3层，如果是返回[4.0, 1.0, 0.4], 否则返回[4, 1, 0.25, 0.06, 0.02]  这些value值是给小中大目标Head给的权重
        self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02])  # P3-P7
        self.ssi = list(m.stride).index(16) if autobalance else 0  # stride 16 index 返回步长为16的索引
        self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
        self.na = m.na  # number of anchors
        self.nc = m.nc  # number of classes
        self.nl = m.nl  # number of layers
        self.anchors = m.anchors
        self.device = device

样本匹配--build_targets函数

样本的匹配需要看build_targets函数。函数中的p是三个预测层，targets数据集的真实值。

tcls,tbox,indices,anch是用来存放匹配结果。

gain是用来后面target缩放到特征层上。

    def build_targets(self, p, targets):
        # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
        na, nt = self.na, targets.shape[0]  # number of anchors, targets.  anchor数量， target数量
        tcls, tbox, indices, anch = [], [], [], []
        gain = torch.ones(7, device=self.device)  # normalized to gridspace gain

ai可以理解为anchor的索引，我们知道yolov5有3个head，每个head上3种anchor，因此这里就相当于给每个head上的anchor编号为0,1,2.通过view将其以列的形式排列，利用repeat函数复制nt列[这里的nt就是target中的目标数].

ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt)

ai tensor([[0., 0., 0.],
           [1., 1., 1.],
           [2., 2., 2.],
        ...,
                       ])
行数等于target数量,shape is [target[0], 1]

通过将targets和前面的ai进行cat拼接操作，那么就可以相当于给每个目标都分配了anchor的索引。

targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2)  # append anchor indices

可以看到给target中的三个目标均分配了0,1,2的索引【后面会利用这个索引来具体看样本和哪个anchor是匹配的】

给每个target分配anchor索引

tensor([[[0.00000, 0.00000, 0.04204, 0.21125, 0.08408, 0.36503, 0.00000],
[0.00000, 0.00000, 0.14960, 0.24400, 0.23867, 0.36503, 0.00000],
[0.00000, 0.00000, 0.36253, 0.24517, 0.21995, 0.39545, 0.00000]],

    [   [0.00000, 0.00000, 0.04204, 0.21125, 0.08408, 0.36503, 1.00000],
         [0.00000, 0.00000, 0.14960, 0.24400, 0.23867, 0.36503, 1.00000],
         [0.00000, 0.00000, 0.36253, 0.24517, 0.21995, 0.39545, 1.00000]],

    [   [0.00000, 0.00000, 0.04204, 0.21125, 0.08408, 0.36503, 2.00000],
         [0.00000, 0.00000, 0.14960, 0.24400, 0.23867, 0.36503, 2.00000],
         [0.00000, 0.00000, 0.36253, 0.24517, 0.21995, 0.39545, 2.00000]]])

targets:[num_obj, 6],repeat(3,1,1)表示复制3个1行1列,则repeat后 targets shape变为[3,num_obj,6]

ai = [[[0.],
      [0.],
      [0.]..]
      
      [[1.],
      [1.],
      [1.]...]
      
      [[2.],
      [2.]
      [2.]...]]
      ai[...,None] shape [3, num_obj, 1]
cat后targets shape [3, num_obj,6+1]=[3,num_obj,7]
也就是targets[...,:6]保存的是targets信息，targets[...,6:]保存的是对于anchors索引

接下来是遍历三个head进行target和anchor的匹配来确定正样本。

比如在遍历第一个head的时候shape为【batch,3,80,80,5+num_classes】.

那么此时的gain通过下面操作将变为【1,1,80,80,80,80,1】。

        for i in range(self.nl):
            anchors, shape = self.anchors[i], p[i].shape
            # torch.tensor(shape)=[batch,3,80,80,85]
            gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]]  # xyxy gain  获取特征层的w和h[80,80,80,80]

大家可以回看target中的box信息，可以看到这些最初的box信息值范围是0~1的，但此时我们的特征层head w和h是80*80，这肯定是不匹配的，所以可以通过targets * gain将这些Box缩放到特征层上得到真实的尺寸。

# Match targets to anchors  先验框和target框的匹配
t = targets * gain  # shape(3,n,7)  将targets中的box缩放到特征层上

那么此时得到的t如下。可以看到box信息缩放到80*80特征层上具体是多少了。

tensor([[[ 0.00000, 0.00000, 3.36320, 16.90000, 6.72640, 29.20240, 0.00000],
[ 0.00000, 0.00000, 11.96800, 19.52000, 19.09360, 29.20240, 0.00000],
[ 0.00000, 0.00000, 29.00240, 19.61360, 17.59600, 31.63600, 0.00000]],

    [[ 0.00000,  0.00000,  3.36320, 16.90000,  6.72640, 29.20240,  1.00000],
      [ 0.00000,  0.00000, 11.96800, 19.52000, 19.09360, 29.20240,  1.00000],
      [ 0.00000,  0.00000, 29.00240, 19.61360, 17.59600, 31.63600,  1.00000]],

    [[ 0.00000,  0.00000,  3.36320, 16.90000,  6.72640, 29.20240,  2.00000],
      [ 0.00000,  0.00000, 11.96800, 19.52000, 19.09360, 29.20240,  2.00000],
      [ 0.00000,  0.00000, 29.00240, 19.61360, 17.59600, 31.63600,  2.00000]]])

yolov5中的样本匹配与之前yolov4或者SSD的样本匹配不同，yolov5采用的是宽高比例的匹配策略，不同于iou匹配。

具体做法是：(1)target的宽高与anchor的宽高比得到ratio1【对应代码中的r】

                  (2)anchor的宽高与target的宽高比得到ratio1【对应代码中的1/r】

                  (3)取上面两个比值的最大值与阈值比较【默认为4对应代码中anchor_t】，如果**小于该阈值那么认为匹配成功**。

r = t[..., 4:6] / anchors[:, None]  # wh ratio
j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t']  # compare

此时得到的j【也就是和阈值相比较】均为False,就是表示没有匹配成功，也就是可以认为，此时的target中所有的目标不能被80*80的head所预测【或者说这些目标中没有小目标】

tensor([[False, False, False],
[False, False, False],
[False, False, False]])

当遍历到第二个head，也就是大小为4040的head上时。target也需要缩放到4040上面，此时为：

tensor([[[ 0.00000, 0.00000, 1.68160, 8.45000, 3.36320, 14.60120, 0.00000],
[ 0.00000, 0.00000, 5.98400, 9.76000, 9.54680, 14.60120, 0.00000],
[ 0.00000, 0.00000, 14.50120, 9.80680, 8.79800, 15.81800, 0.00000]],

    [    [ 0.00000,  0.00000,  1.68160,  8.45000,  3.36320, 14.60120,  1.00000],
          [ 0.00000,  0.00000,  5.98400,  9.76000,  9.54680, 14.60120,  1.00000],
          [ 0.00000,  0.00000, 14.50120,  9.80680,  8.79800, 15.81800,  1.00000]],

    [    [ 0.00000,  0.00000,  1.68160,  8.45000,  3.36320, 14.60120,  2.00000],
          [ 0.00000,  0.00000,  5.98400,  9.76000,  9.54680, 14.60120,  2.00000],
          [ 0.00000,  0.00000, 14.50120,  9.80680,  8.79800, 15.81800,  2.00000]]])

此时通过宽高比阈值筛选得到：也就是可以知道target中目标可以被40*40这个head的第0与第2号anchor所匹配，这些anchor所匹配的样本就是我们要的正样本GT.

tensor([[ True, False, False],
[False, False, False],
[ True, True, True]])

通过上面的mask过滤一下targets.

tensor([[ 0.00000, 0.00000, 1.68160, 8.45000, 3.36320, 14.60120, 0.00000],
[ 0.00000, 0.00000, 1.68160, 8.45000, 3.36320, 14.60120, 2.00000],
[ 0.00000, 0.00000, 5.98400, 9.76000, 9.54680, 14.60120, 2.00000],
[ 0.00000, 0.00000, 14.50120, 9.80680, 8.79800, 15.81800, 2.00000]])

得到过滤后目标的中心点坐标

gxy = t[:, 2:4]  # grid xy

得到中心点相对于边界的距离

gxi = gain[[2, 3]] - gxy

jk和lm是判断gxy的中心点更偏向哪里.

j, k = ((gxy % 1 < g) & (gxy > 1)).T
l, m = ((gxi % 1 < g) & (gxi > 1)).T
j = torch.stack((torch.ones_like(j), j, k, l, m))

    得到的**j**如下，包含当前网格有五个cell，第一行保留所有的gtbox，第二行表示左边的cell中的gt，第三行是表示上方的cell中的gt，第四行是右边cell的网格，第五行是下方的cell中的gt。这里与v3和v4不同在于之前的yolo是目标落在哪个head的cell就由该cell进行预测，而v5通过增加邻近的cell来预测，这样就是相当于**增加了正样本的数量**。

tensor([[ True, True, True, True],
[False, False, False, False],
[ True, True, False, False],
[ True, True, True, True],
[False, False, True, True]])

在yolov5中不仅仅用了中心点进行预测，还采用了距离中心点网格最近的两个网格，所以是有五种情况【四周的网格和当前中心的网格】同时用上面的j过滤，这样就可以得出哪些网格有目标

t = t.repeat((5, 1, 1))[j]
offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]

    def build_targets(self, p, targets):
        # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
        na, nt = self.na, targets.shape[0]  # number of anchors, targets.  anchor数量， target数量
        tcls, tbox, indices, anch = [], [], [], []
        gain = torch.ones(7, device=self.device)  # normalized to gridspace gain
        '''
        ai tensor([[0., 0., 0.],
                [1., 1., 1.],
                [2., 2., 2.],
                ...,
                ])
        行数等于target数量,shape is [target[0], 1]
        '''
        ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
        '''
        targets:[num_obj, 6],repeat(3,1,1)表示复制3个1行1列,则repeat后 targets shape变为[3,num_obj,6]
        ai = [[[0.],
              [0.],
              [0.]..]
              
              [[1.],
              [1.],
              [1.]...]
              
              [[2.],
              [2.]
              [2.]...]]
              ai[...,None] shape [3, num_obj, 1]
        cat后targets shape [3, num_obj,6+1]=[3,num_obj,7]
        也就是targets[...,:6]保存的是targets信息，targets[...,6:]保存的是对于anchors索引
        '''
        targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2)  # append anchor indices

        g = 0.5  # bias
        off = torch.tensor(
            [
                [0, 0],
                [1, 0],
                [0, 1],
                [-1, 0],
                [0, -1],  # j,k,l,m
                # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
            ],
            device=self.device).float() * g  # offsets

        for i in range(self.nl):
            anchors, shape = self.anchors[i], p[i].shape
            # torch.tensor(shape)=[batch,3,80,80,85]
            gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]]  # xyxy gain  获取特征层的w和h[80,80,80,80]

            # Match targets to anchors  先验框和target框的匹配
            t = targets * gain  # shape(3,n,7)  将targets中的box缩放到特征层上
            if nt:
                # Matches
                '''
                yolov5采用宽高比例的匹配策略，不同于iou匹配。
                target的宽高与anchors宽高对应相除得到ratio1
                anchors与target的宽高相处得到ratio2[也就是代码中的1/r]
                取两个ratio最大值作为最后的宽高比，该宽高比和设定的阈值(默认为4[anchor_t])比较，小于该阈值的anchor则为匹配到的anchor
                '''
                r = t[..., 4:6] / anchors[:, None]  # wh ratio
                j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t']  # compare
                # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
                t = t[j]  # filter

                # Offsets
                gxy = t[:, 2:4]  # grid xy
                gxi = gain[[2, 3]] - gxy  # inverse
                j, k = ((gxy % 1 < g) & (gxy > 1)).T
                l, m = ((gxi % 1 < g) & (gxi > 1)).T
                j = torch.stack((torch.ones_like(j), j, k, l, m))
                t = t.repeat((5, 1, 1))[j]
                offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
            else:
                t = targets[0]
                offsets = 0

            # Define
            bc, gxy, gwh, a = t.chunk(4, 1)  # (image, class), grid xy, grid wh, anchors
            a, (b, c) = a.long().view(-1), bc.long().T  # anchors, image, class
            gij = (gxy - offsets).long()
            gi, gj = gij.T  # grid indices

            # Append
            indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1)))  # image, anchor, grid
            tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
            anch.append(anchors[a])  # anchors
            tcls.append(c)  # class

        return tcls, tbox, indices, anch

通过上面的样本匹配操作我们会得到4个值。

tcls:存储target中的类别，tbox：gt中的box信息，indices：当前gtbox属于第几张图像，gtbox与anchor的对应关系以及所属的cell坐标。anchors:anchor信息

获得b：图像；a:anchor, gj:gi Cell的纵坐标与横坐标

b, a, gj, gi = indices[i]

tobj是用来后面存储gt中的目标信息，shape[batch_size,3, 80,80]

tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device)

如果目标存在，获取当前head所预测的中心坐标pxy，pwh,以及类置信度

pxy, pwh, _, pcls = pi[b, a, gj, gi].split((2, 2, 1, self.nc), 1)  # target-subset of predictions

box_loss

利用iou获取box_loss。

                # Regression
                pxy = pxy.sigmoid() * 2 - 0.5
                pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
                pbox = torch.cat((pxy, pwh), 1)  # predicted box
                iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze()  # iou(prediction, target)
                lbox += (1.0 - iou).mean()  # iou loss

cls_loss

                if self.nc > 1:  # cls loss (only if multiple classes)
                    t = torch.full_like(pcls, self.cn, device=self.device)  # targets
                    t[range(n), tcls[i]] = self.cp
                    lcls += self.BCEcls(pcls, t)  # BCE

obj_loss

pi是指当前head层,pi[...,4]即取出conf这一维度与tobj做交叉熵。

obji = self.BCEobj(pi[..., 4], tobj)

再与所对应的权重相乘，得到obj_loss

lobj += obji * self.balance[i]

最后返回Loss

        lbox *= self.hyp['box']
        lobj *= self.hyp['obj']
        lcls *= self.hyp['cls']
        bs = tobj.shape[0]  # batch size

        return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach()

loss: tensor([10.73431]) loss_item: tensor([0.06105, 5.30610, 0.00000])