yolov5 loss函数理解
最近整理了一下yolov5 loss的函数，关键点都在代码段里
YOLOv5的loss主要由三个部分组成：
1、Classes loss，分类损失，采用BCE loss，只计算正样本的分类损失。
2、Objectness loss，obj置信度损失，采用BCE loss，计算的是所有样本的obj损失。注意这里的obj指的是网络预测的目标边界框与GT Box的CIoU。
3、Location loss，定位损失，采用CIoU loss，只计算正样本的定位损失。
class ComputeLoss:
# Compute losses
def __init__(self, model, autobalance=False):
super(ComputeLoss, self).__init__()
# 获取模型在cpu还是gpu上运行的.之后生成的临时变量也会在相应的设备上运行
device = next(model.parameters()).device # get model device
h = model.hyp # hyperparameters
# Define criteria
# 定义cls loss和 obj loss,对象实例化
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)
# 获取模型的detect层
m = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module
# 用来实现obj,box,cls loss在每一层之间权重的平衡
self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl,
[4.0, 1.0, 0.25, 0.06, .02]) # P3-P7 的特征点的数量比值为P3 : P4 : P5 = 4:1:0.25
# 获取各个特征层的stride相关参数
self.ssi = list(m.stride).index(16) if autobalance else 0 # stride 16 index
# 将各个loss加入到类的公共变量中
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, 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
def __call__(self, p, targets): # predictions, targets, model
# loss计算
# p为模型的预测输出 p的最后一维cx,cy,w,h,conf + number of class(coco == 80)
# targets为gt框的信息
device = targets.device
# 初始化loss
lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
# 建立targets目标,获取扩充后的正样本
# tcls [[num]] 存放了gt框所对应的网格的cls
# tbox [[x_offset,y_offset,w,h]] 存放了gt框所对应的网格的box,注意此处的x和y是相对于网格的偏移量offset
# 这个offset相对于indices [[image, anchor, grid indices_y, grid indices_x]]中grid indices的偏移量
# indices [[image, anchor, grid indices(gj, gi)]] 存放了gt对应的gird的信息
# indices 的len==3,对应３层输出，每一层包括:image对应batchsize的哪张图片,anchor,对应哪个尺度的anchor,以及所在的网格
# anchors [[num,2]]#anchor信息
tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets
# Losses 3个feature map上的输出，一层一层来处理
for i, pi in enumerate(p): # layer index, layer predictions
# 获取扩充正样本的targets的图片号，anchor序号，网格的位置
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
# obj pi.shape=torch.size([bs,na,featuresize_h,featuresize_w,7])
# pi[..., 0]获取最后一维的第一维，最后一维表示的是每一个目标的７个标签，所以这里的维数表示的即是当前输出层的targets的总数
# tobj.shape=torch.size([bs,na,featuresize_h,featuresize_w])
# 注意这里tobj的size不只是正样本的shape,而是所有网格的所有anchor对应的shape,
# 即除了正样本还包括built根据标签过滤掉的负样本，以及标签之外背景处的anchor
tobj = torch.zeros_like(pi[..., 0], device=device) # target
n = b.shape[0] # number of targets 正样本总数
if n:
# 对应targets的预测值，pi是每一层feature map的预测tensor
# pi.shape = (bs, na, feature_w, feature_h, 4+1+class_num)
# gj,gi是中心点所在feature map位置，是采用targets中心点所在位置的anchor来回归。
# 获取和正样本targets相同图片名，相同anchor, 相同网格位置的预测框的标签信息
# ps.shape = (nt, 4+1+class_num) cx,cy,w,h,conf + number of class
ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
# Regression 目标框回归
# 因为在build_targets中，作者相当于扩充了标签框所在网格的上下左右４个网格，
# 左上网格的中心点偏移了－0.5，右下网格偏移了0.5，所以在中心坐标回归到的时候限制了回归的范围到偏移的网格范围内。
# yolov4中sigmoid(tx)+cx的范围是cx-1,cx+1,当tx=0时，cx+0.5，认为中心点偏移的范围在相邻两个网格内
# 当前方法范围是1.5+cx, cx-0.5, 当tx=0时，cx=0.5+cx，认为中心点坐标的偏移的最大范围就是扩充正样本时的范围，
# 左上角网格-0.5,右下角网格+0.5,所以总的范围就是-0.5 ~ 1.5
pxy = ps[:, :2].sigmoid() * 2. - 0.5
# w,h 回归其没有采用exp操作，而是直接乘上anchors[i]。
# yolov4中求回归框的长和高的时候，直接对tw做指数操作保证缩放的系数大于0,但是宽度和高度的最大值完全不受限制，
# 这种指数的运算很危险，因为它可能导致失控的梯度、不稳定、NaN损失并最终完全失去训练。
# 所以yolov5对w,h 回归其没有采用exp操作，而是用sigmoid来限制了缩放系数的最大值。因为作者在YOLO5有一个超参数为anchor_t，就是4。
# 该超参数的使用方法是，在训练时，如果真实框与锚框尺寸的比值大于4，限于预测框回归公式上界，该锚框是没有办法与该真实框重合的，
# 所以训练时会把比值大于4的锚框删除掉。
# 作者认为回归框和anchor的最大比值是４，所以将缩放系数的最大值设为４=2**2。
pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
# 找到和targets中相同图片相同anchor,相同网格的预测框，回归后的集合pbox。所以和正样本不在同一个网格的预测框并不用回归坐标
pbox = torch.cat((pxy, pwh), 1) # predicted box
# 计算ciou
iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target)
# 计算box的ciouloss
lbox += (1.0 - iou).mean() # iou loss
# Objectness
# 获取target所对应的obj,网格中存在gt目标的会被标记为iou与gt的交并比 gr定义在train.py model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou)
# 神经网络的训练有时候可能希望保持一部分的网络参数不变，只对其中一部分的参数进行调整；或者只训练部分分支网络，并不让其梯度对主网络的梯度造成影响，
# torch.tensor.detach()和torch.tensor.detach_()函数来切断一些分支的反向传播
# torch.clamp(input, min, max, out=None) 限幅。将input的值限制在[min, max]之间，并返回结果
# torch.tensor.type该方法的功能是:当不指定dtype时,返回类型.当指定dtype时,返回类型转换后的数据,如果类型已经符合要求,那么不做额外的复制,返回原对象.
# iou.detach().clamp(0)将iou中所有的小于0的置0，认为小于0的CIOU是负样本
# 给正样本的tobj赋初值，初值里用到了iou取代1，代表该点对应置信度，负样本，包括背景的置信度为0
# 引入了大量正样本anchor，但是不同anchor和gt bbox匹配度是不一样，预测框和gt bbox 的匹配度也不一样，
# 如果权重设置一样肯定不是最优的，故作者将预测框和bbox的iou作为权重乘到conf分支，用于表征预测质量。
# 一般检测网络的分类头，在计算loss阶段，标签往往是非0即1的状态，即是否为当前类别。
# yolov5则是将anchor与目标匹配时的iou作为该位置样本的标签值。iou值在0-1之间，label值的缩小导致了最后预测结果值偏小。
# 通过model.gr可以修改iou值所占权重，默认是1.0，即用iou值完全作为标签值,而不是非0即1。iou的最大值为1．
tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio
# Classification
if self.nc > 1: # cls loss (only if multiple classes)
# t.shape = torch.size(nt,nc) t为nt*nc的0
t = torch.full_like(ps[:, 5:], self.cn, device=device) # targets
# 在nt*nc中，第1维数据取所有，即range(n),n第i层feature map输出的正样本的数据,第2维数据取第i层输出数据的标签tcls[i]中的类为1,其余为0
# 每一个target对应nc类，所以t的第2个维度是nc, 这里就是将第n个target对应的类别置为cp＝1，其余为0
# 加入第一个数据的类别是2，共有4类，则t[0]=[0,0,1,0],简单理解就是这里做了个one hot embedding
t[range(n), tcls[i]] = self.cp
# 预测值和真值做二进制交叉熵，但是只计算了带有GT的预测值，不带预测值的不计算,
# 计算时每一类认为类别是1，其余类别是0，不带背景计算，参与计算的是所有的类
# bceloss中最终返回的是平均每张图片的每个类别的平均损失
lcls += self.BCEcls(ps[:, 5:], t) # BCE
# 将整个数据中给每个cell的每个anchor是否对应目标都用于计算loss,即不仅考虑了正样本，还考虑了背景
# 由以上可知lbos,lcls，都没有考虑背景，只有lobj考虑了背景
obji = self.BCEobj(pi[..., 4], tobj)
# self.balance[i]对每一层输出的lobj都加了权重,
# 3个预测分支上的分类损失、坐标损失直接相加，得到总的分类损失和坐标损失，
# 而3个预测分支上的目标置信度损失需要进行加权再相加，得到总的目标置信度损失，
# 权值分别为[4.0, 1.0, 0.4]，其中4.0是用在大特征图（预测小目标）上，
# 所以，这里的加权，我认为是旨在提高小目标的检测精度。
lobj += obji * self.balance[i] # obj loss
if self.autobalance:
self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
if self.autobalance:
self.balance = [x / self.balance[self.ssi] for x in self.balance]
# 对lbox，lobj，lcls增加权重
lbox *= self.hyp['box']
lobj *= self.hyp['obj']
lcls *= self.hyp['cls']
bs = tobj.shape[0] # batch size
# 最后将分类损失、坐标损失、置信度损失直接相加，得到一个总损失（一个batch中每张图像的平均总损失），再乘以batch的大小，得到用于更新梯度的损失。
loss = lbox + lobj + lcls
return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
def build_targets(self, p, targets):
# Build targets for compute_loss(), input targets(image,class,x,y,w,h);targets就是标注的gt框
na, nt = self.na, targets.shape[0] # number of anchors(each layer), targets(每个batch中的标签个数)
# 初始化每个batch box的信息 tcls表示类别，tbox表示标记的box和生成的box的坐标（x,y,w,h），indices表示图像索引，anch表示选取的anchor的索引
tcls, tbox, indices, anch = [], [], [], []
gain = torch.ones(7, device=targets.device) # normalized to gridspace gain 将targets
# torch.arange(start, end)创建从start到end的int64的tensor。
# repeat()函数可以对张量进行重复扩充,当参数只有两个时：（列的重复倍数，行的重复倍数）。1表示不重复,所以这里列不重复，行重复nt次
# 此时得到的ai.shape = [3, nt]
ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)
# 原本targets.shape=[nt,6],[image,class,x,y,w,h]，repeat根据每一层的anchor数量将targets增加一维,shape[na=3,nt,6]
# ai[:, :, None] 2维变３维torch.Size([3, nt, 1])
# torch.cat两个3维的tensort在第2维上concat,targets的torch.Size([3, nt, 7])
# 第一维增加layer的索引，并且将原本targets[image,class,x,y,w,h]在最后上增加anchor的索引[image,class,x,y,w,h,anchor indices],
# 也就说把每个gt框分配给了每一层输出的每一个anchor
targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices
g = 0.5 # bias 什么用？ 为扩充标记的box添加偏置，具体扩充规则为在下边,目的为了增加正样本
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 偏置矩阵
# nl预测头的数量，输出layer的数量。顺序为降采样8-16-32
# anchor匹配需要逐层进行。不同的预测层其特征图的尺寸不一样，而targets是相对于输入分辨率的宽高做了归一化，
# targets * gain通过将归一化的box坐标投影到特征图上。
for i in range(self.nl):
anchors = self.anchors[i] # 获取该层特征图中的anchor shape=[３，２】
gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain 获取该层特征图的尺寸（1,1，w,h,w,h,1）
# Match targets to anchors
# #targets[image,class,x,y,w,h,anchor indices]的box坐标转换到特征图上
# 在box坐标生成的时候，对box坐标进行了归一化，即除以图像的宽高），通过将归一化的box乘以特征图的尺度，
# 从而将坐标投影到特征图上,此时t和targets尺寸一样，[3,nt,7]
t = targets * gain
# Matches 当该batch中存在标签，获取每个box对应的anchor，并生成对符合规定的anchor
if nt:
# yolov5抛弃了MaxIOU匹配规则而采用shape匹配规则，计算标签box和当前层的anchors的宽高比，即:wb/wa,hb/ha。
# 如果宽高比大于设定的阈值说明该box没有合适的anchor，在该预测层之间将这些box当背景过滤掉。
# 对b中保存下来的gt进行扩充:
# 1)保存现有所有的gt
# 2)保存box中心点坐标Xc距离网格左边的距离小于0.5且坐标大于1的box
# 3)保存box中心点坐标Yc距离网格上边的距离小于0.5且坐标大于1的box
# 4)保存box中心点坐标Xc距离网格右边的距离小于0.5且坐标大于1的box
# 5)保存box中心点坐标Yc距离网格下边的距离小于0.5且坐标大于1的box 将该5中box构成需要的gt
#
# 补充:为什么会取距离四边小于0.5的点,是因为等于0.5时,我们认为该box正好落到该网格中,但是小于0.5时,
# 可能是因为在网络不断降采样时,对特征图尺度进行取整导致box中心产生了偏差,
# 所以作者将小于0.5的box减去偏置1（off矩阵）,使得box中心移动到相邻的特征图网格中,
# 从而对正样本进行扩充,保证了偏差导致的box错位以及扩充了正样本的个数
# 获取box和3个anchor的对应的的长宽比
r = t[:, :, 4:6] / anchors[:, None] # wh ratio r.shape=[3,nt,2]
# 如果长宽比中的最大值小于anchor_t，则该box是合适的box，并获取对应的anchor
j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare j.shape=[3,nt]
# 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 t.shape=[182,7] 按照该匹配策略,一个gt box可能在同一层同时匹配上多个anchor。
# Offsets 获取选择完成的box的*中心点*坐标-gxy（以图像左上角为坐标原点），并转换为以特征图右下角为坐标原点的坐标-gxi
gxy = t[:, 2:4] # grid xy
gxi = gain[[2, 3]] - gxy # inverse
# 分别判断box的（x，y）坐标是否大于1，并距离网格左上角的距离（准确的说是y距离网格上边或x距离网格左边的距离）小于0.5，
# 如果（x，y）中满足上述两个条件，则选中.gxy.shape=[182,2]，包含x,y,所以判别后转置得到j,k,2个结果
# 对转换之后的box的（x，y）坐标分别进行判断是否大于1，并距离网格右下角的距离（准确的说是y距离网格下边或x距离网格右边的距离）距离小于0.5，
# 如果（x，y）中满足上述两个条件，为Ture
j, k = ((gxy % 1. < g) & (gxy > 1.)).T
l, m = ((gxi % 1. < g) & (gxi > 1.)).T
# 获取所有符合要求的box,将原始box和扩增的box进行合并
j = torch.stack((torch.ones_like(j), j, k, l, m))
t = t.repeat((5, 1, 1))[j]
# 生成所有box对应的偏置
offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
else:
t = targets[0]
offsets = 0
# Define
# 获取保留的gt中中心点坐标和网格点坐标的偏执与box的wh构成新的box信息(Cx-X,Cy-Y,w,h)
# 最后获取4个向量
# a.indice[图像序号,anchor序号,网格点坐标x,网格点坐标y]
# b.tbox box的对应坐标[Cx-X,Cy-Y,w,h]
# c.anch tbox对应的anchor索引
# d.tcls tbox对应的类别
# 获取每个box的图像索引和类别
b, c = t[:, :2].long().T # image, class
# 获取box的xy和wh
gxy = t[:, 2:4] # grid xy
gwh = t[:, 4:6] # grid wh
# 获取每个box所在的网格点坐标
gij = (gxy - offsets).long()
gi, gj = gij.T # grid xy indices
# Append
# 获取每个anchor索引
a = t[:, 6].long() # anchor indices
# 保存图像序号 anchor序号和网格点坐标
indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices
# 获取（x,y)相对于网格点的偏置，以及box的宽高
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
anch.append(anchors[a]) # anchors
tcls.append(c) # class
return tcls, tbox, indices, anch