在前两篇博客中我分析了YOLO v7的模型结构以及训练数据的准备。这里将对损失函数的代码进行分析。

在train.py中,我们可以看到以下的代码是进行损失值计算的

if 'loss_ota' not in hyp or hyp['loss_ota'] == 1:loss, loss_items = compute_loss_ota(pred, targets.to(device), imgs)# loss scaled by batch_sizebreakelse:loss, loss_items = compute_loss(pred, targets.to(device))# loss scaled by batch_size

因为在超参数中定义了loss_ota,因此是用compute_loss_ota来进行计算。这个对应的是loss.py里面定义的ComputeLossOTA的类。

查看这个类里面的__call__函数,其输入变量是模型输出的预测值p, 训练图像包含的targets, 以及训练图像imgs。

例如以下是输入的一幅训练图片:

图片包含的targets如下:

tensor([[0.0000e+00, 4.5000e+01, 5.1015e-01, 8.6860e-02, 9.6169e-01, 1.7279e-01],[0.0000e+00, 5.0000e+01, 8.6156e-01, 8.8347e-02, 2.5677e-01, 1.7359e-01],[0.0000e+00, 1.5000e+01, 9.0563e-01, 8.2660e-01, 1.8675e-01, 3.4358e-01],[0.0000e+00, 3.9000e+01, 7.5925e-02, 1.2316e-01, 1.5037e-01, 2.3427e-01],[0.0000e+00, 8.0000e+00, 5.1819e-01, 1.0259e-01, 4.6240e-01, 2.0505e-01],[0.0000e+00, 7.0000e+00, 7.8777e-01, 4.4500e-02, 2.7215e-02, 2.3584e-02],[0.0000e+00, 6.0000e+00, 5.3230e-01, 4.1727e-01, 5.3604e-01, 3.4526e-01],[0.0000e+00, 1.6000e+01, 5.0707e-01, 3.8664e-01, 3.7291e-02, 3.2220e-02]], device='cuda:0')

这个targets数组共包括了8个target,每个target对应一个6维的变量,第一维代表Batch的第几个图片,第2维代表物体的类别,后面4维对应中心点的xy以及wh,这里xywh是表示target的大小与图片的大小比例。

模型输出的预测值P是一个列表,包含大中小三个维度的预测结果,其维度分别为(batch_size, 3, 80, 80, 85), (batch_size, 3, 40, 40, 85), (batch_size, 3, 20, 20, 85)。这里面3表示有3个anchor,80*80代表最细粒度的检测,把图片分为80*80个格子,85表示每个格子的一个预测变量,其中包括了预测物体存在概率,物体的xywh以及80个分类的概率,总共是85维。

我们把targets数组的物体在图像中标识出来看一下:

在计算loss之前,首先对输入的targets进行处理。在ComputeLossOTA里面有一个build_targets,对原始的targets进行处理,生成更多的正样本。

首先是对于每个target,根据不同的anchor和邻近的网格,初步筛选生成候选正样本,如以下的代码:

indices, anch = self.find_3_positive(p, targets)

这个函数的定义如下,我添加了注释:

def find_3_positive(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, targetsindices, anch = [], []gain = torch.ones(7, device=targets.device).long()# normalized to gridspace gainai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)# same as .repeat_interleave(nt)#把target的第一个维度重复na次,并与anchor的编号拼接#例如原来target的维度为[8,6],处理后的维度为[3,8,7]targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)# append anchor indicesg = 0.5# bias#off表示包括了不同的偏移量,分别对应target本身所处的网格以及上下左右4个邻近的网格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=targets.device).float() * g# offsets#分别基于三个检测粒度对target进行处理for i in range(self.nl):anchors = self.anchors[i]#以最小的检测粒度80*80网格为例,gain是[1, 1, 80, 80, 80, 80,1]gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]# xyxy gain# Match targets to anchors,即把xywh和网格对应起来t = targets * gainif nt:# Matchesr = t[:, :, 4:6] / anchors[:, None]# wh ratio,计算每个target的宽高与anchor的比值j = torch.max(r, 1. / r).max(2)[0]  model.hyp['iou_t']# iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))t = t[j]# filter# Offsetsgxy = t[:, 2:4]# grid xy,获取target对应的xy的网格坐标 # 计算target的附近网格是否也作为正样本。例如某个target的xy坐标为(63.23, 21.75),那么(62,22)这个坐标的网格也是这个target的一个正样本gxi = gain[[2, 3]] - gxy# inverse,计算网格数与target对应的xy的网格坐标的差值j, k = ((gxy % 1.  1.)).T#判断xy网格坐标如果是大于1并且其余数是小于0.5,那么代表其左边或上边的网格也负责这个target的预测l, m = ((gxi % 1.  1.)).T#同上,代表右边或下边的网格也负责这个target的预测j = torch.stack((torch.ones_like(j), j, k, l, m)) #拼接之后的维度为[5, x],其中5表示分别对应自身网格,以及上下左右四个网格的判断结果。x表各个anchor对应的总样本数。列如初始的target总数为8,其中有2个target与80*80这个粒度的3个anchor的宽高比值的最大值都小于4,因此x=2*3=6t = t.repeat((5, 1, 1))[j]offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]#offset为偏移量else:t = targets[0]offsets = 0# Defineb, c = t[:, :2].long().T# image, classgxy = t[:, 2:4]# grid xygwh = t[:, 4:6]# grid whgij = (gxy - offsets).long()#根据偏移量计算正样本对应的网格gi, gj = gij.T# grid xy indices# Appenda = t[:, 6].long()# anchor indicesindices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1)))# image, anchor, grid indicesanch.append(anchors[a])# anchorsreturn indices, anch

还是以之前举例的targets为例,当检测粒度是80*80网格的时候,这8个target当中只有两个,其wh与80*80的三个anchor的比值的最大值都小于4,即以下两个target,分别对应上图的truck和dog:

tensor([[0.0000e+00, 7.0000e+00, 7.8777e-01, 4.4500e-02, 2.7215e-02, 2.3584e-02],[0.0000e+00, 1.6000e+01, 5.0707e-01, 3.8664e-01, 3.7291e-02, 3.2220e-02]], device='cuda:0')

这两个target的xy网格坐标分别为(63.0125, 3.56),(40.5656, 30.9312)。因此其周围的网格(62, 3), (63, 4), (41, 30), (40, 31)也是正样本,加上自身所处的网格,总共是6个网格。总的样本数为3*6=18

得到这些候选的正样本之后,下一步是对这些正样本进行进一步筛选。

把这些正样本对应的网格的预测值取出,计算其与每个target之间的IOU

for i, pi in enumerate(p): #p是三个维度的预测值的列表,p[0]的维度是[3,80,80,85]#获取之前候选正样本的值,b表示batch里面的某一个,a表示anchor,gj,gi是格子的坐标b, a, gj, gi = indices[i]idx = (b == batch_idx)b, a, gj, gi = b[idx], a[idx], gj[idx], gi[idx]all_b.append(b)all_a.append(a)all_gj.append(gj)all_gi.append(gi)all_anch.append(anch[i][idx])from_which_layer.append(torch.ones(size=(len(b),)) * i)#获取正样本对应的预测值fg_pred = pi[b, a, gj, gi]p_obj.append(fg_pred[:, 4:5])p_cls.append(fg_pred[:, 5:])grid = torch.stack([gi, gj], dim=1)#计算xy,wh的预测值,注意这里要对sigmoid的值*2-0.5,因为在前面初筛正样本里面,我们除了target所处的格子,其四周的格子也会负责预测,因此这个xy的值对sigmoid的值*2-0.5后,其范围是[-0.5, 1.5]pxy = (fg_pred[:, :2].sigmoid() * 2. - 0.5 + grid) * self.stride[i] #/ 8.#pxy = (fg_pred[:, :2].sigmoid() * 3. - 1. + grid) * self.stride[i]pwh = (fg_pred[:, 2:4].sigmoid() * 2) ** 2 * anch[i][idx] * self.stride[i] #/ 8.pxywh = torch.cat([pxy, pwh], dim=-1)pxyxy = xywh2xyxy(pxywh)pxyxys.append(pxyxy)pxyxys = torch.cat(pxyxys, dim=0)if pxyxys.shape[0] == 0:continuep_obj = torch.cat(p_obj, dim=0)p_cls = torch.cat(p_cls, dim=0)from_which_layer = torch.cat(from_which_layer, dim=0)all_b = torch.cat(all_b, dim=0)all_a = torch.cat(all_a, dim=0)all_gj = torch.cat(all_gj, dim=0)all_gi = torch.cat(all_gi, dim=0)all_anch = torch.cat(all_anch, dim=0)#计算预测值与target的iou。例如我们之前是初筛获得了99个正样本,target是8,那么pair_wise_iou的维度是[8,99],即每个target都分别与这99个正样本计算ioupair_wise_iou = box_iou(txyxy, pxyxys)#计算iou_loss,如果iou越大,则这个loss越小pair_wise_iou_loss = -torch.log(pair_wise_iou + 1e-8)

之后要对这个IOU进行排序,取出其头10位最大的IOU数值(如果不足10个则按实际数量取),然后把这10个IOU相加,结果取整(如结果小于1则取1)。这个值将表示对应某个target的正样本数量。

top_k, _ = torch.topk(pair_wise_iou, min(10, pair_wise_iou.shape[1]), dim=1)dynamic_ks = torch.clamp(top_k.sum(1).int(), min=1)

例如我们计算获得了之前的8个target与这99个正样本的IOU数值,然后取其前10位最大的IOU,数值如下:

tensor([[0.33164, 0.31670, 0.31670, 0.27168, 0.27168, 0.26650, 0.26650, 0.24327, 0.23278, 0.23109],[0.59865, 0.57707, 0.46241, 0.41670, 0.38554, 0.37295, 0.36225, 0.33483, 0.33090, 0.22409],[0.74503, 0.71696, 0.58363, 0.37660, 0.36653, 0.36114, 0.34860, 0.32036, 0.29622, 0.16013],[0.46875, 0.46193, 0.39073, 0.30311, 0.27263, 0.17448, 0.14117, 0.11435, 0.09316, 0.07584],[0.53946, 0.53946, 0.44517, 0.44517, 0.38395, 0.36339, 0.36339, 0.31876, 0.31780, 0.31460],[0.46457, 0.30825, 0.29158, 0.20381, 0.20071, 0.14967, 0.14524, 0.13727, 0.12137, 0.02018],[0.66553, 0.52719, 0.30917, 0.26019, 0.20343, 0.19609, 0.17714, 0.17714, 0.15379, 0.14393],[0.56813, 0.54395, 0.38423, 0.29507, 0.23172, 0.21839, 0.16346, 0.16189, 0.12244, 0.05927]], device='cuda:0', grad_fn=)

经过求和计算之后,dynamic_ks的数值如下:

tensor([2, 4, 4, 2, 4, 2, 2, 2], device='cuda:0', dtype=torch.int32)

所以对于target1来说,其正样本的数量为2,对于target2其正样本的数量为4

计算每个target对应的全部候选正样本的预测分类的loss

gt_cls_per_image = (F.one_hot(this_target[:, 1].to(torch.int64), loss.nc).float().unsqueeze(1).repeat(1, pxyxys.shape[0], 1))num_gt = this_target.shape[0]cls_preds_ = (p_cls.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()* p_obj.unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_())y = cls_preds_.sqrt_()pair_wise_cls_loss = F.binary_cross_entropy_with_logits(torch.log(y/(1-y)) , gt_cls_per_image, reduction="none").sum(-1)

这里计算结果pair_wise_cls_loss的维度是[8,99],把这个loss与之前计算的IOU_loss相加,计算cost

cost = (pair_wise_cls_loss+ 3.0 * pair_wise_iou_loss)

然后就根据cost来进行从小到大的topK的排序,其中的K是刚才计算得到的dynamic_ks,获取K个cost最小值所对应的正样本的编号

matching_matrix = torch.zeros_like(cost)for gt_idx in range(num_gt):_, pos_idx = torch.topk(cost[gt_idx], k=dynamic_ks[gt_idx].item(), largest=False)matching_matrix[gt_idx][pos_idx] = 1.0

例如对于第一个target,dynamic_ks对应的值是2,因此取出2个cost最小的对应的正样本的编号,这里的计算结果pos_idx是

tensor([74, 62], device='cuda:0')

即对于第一个target,我们挑选第74和62号这两个正样本

要注意的是,对于某一个target,我们可以挑选多个正样本来进行训练,也就是多个正样本可以预测同一个target。但是一个正样本不能同时预测多个target,以下代码判断是否有这种情况,如有则取cost最小的那个target作为预测:

anchor_matching_gt = matching_matrix.sum(0)if (anchor_matching_gt > 1).sum() > 0:_, cost_argmin = torch.min(cost[:, anchor_matching_gt > 1], dim=0)matching_matrix[:, anchor_matching_gt > 1] *= 0.0matching_matrix[cost_argmin, anchor_matching_gt > 1] = 1.0

计算这些正样本中哪些被挑选出来预测target的,以及其对应要预测的target的序号

fg_mask_inboxes = matching_matrix.sum(0) > 0.0matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)

把这些挑选出来的正样本所对应的anchor,网格xy坐标,预测的target序号等都进行保存并返回,这些就是筛选后得到的最终的正样本,可以用来进行模型训练的损失函数计算的输入值。

例如上面例子提到的truck这个target,其对应在80*80的细粒度检测层中,对应了[63,3], [62,3], [63,4]这三个正样本,经过进一步筛选后,保留了[63,3], [62,3]这两个正样本。

最后就是损失函数的计算了,这部分比较简单,包括了正样本与target的IOU loss,objectness loss,分类loss这三部分。

IOU loss的计算代码如下:

pxy = ps[:, :2].sigmoid() * 2. - 0.5pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]pbox = torch.cat((pxy, pwh), 1)# predicted boxselected_tbox = targets[i][:, 2:6] * pre_gen_gains[i]selected_tbox[:, :2] -= gridiou = bbox_iou(pbox.T, selected_tbox, x1y1x2y2=False, CIoU=True)# iou(prediction, target)lbox += (1.0 - iou).mean()# iou loss

Objectness loss计算如下,其中对于正样本来说其置信度是IOU的值,对于其他负样本以及背景,其置信度为0:

tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype)obji = self.BCEobj(pi[..., 4], tobj)lobj += obji * self.balance[i]# obj loss

分类loss的计算如下:

selected_tcls = targets[i][:, 1].long()if self.nc > 1:# cls loss (only if multiple classes)t = torch.full_like(ps[:, 5:], self.cn, device=device)# targetst[range(n), selected_tcls] = self.cplcls += self.BCEcls(ps[:, 5:], t)# BCE

最后把这三个loss加起来即可

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

在下一篇博客,我将继续分析Yolov7的训练过程。