yolov5 的损失函数主要是分为三个部分，分别是分类损失、定位损失、置信度损失。
这里我总结出 ComputeLoss 类的流程大致分为5步：
计算出当前模型输出的损失值
核心逻辑 : 分类损失 + 定位损失 + 置信度损失
1. 初始化置信度损失和分类损失所需要用到的BCELoss函数
2. 得到GT(Ground True)正样本相关参数(索引, 分类, 长宽位置等等)
2.1 比较GT和anchor的宽比和高比，大于一定的阈值(4)就是负样本
2.2 筛选到的正样本信息, 观察其中心点坐标, 小于 0.5 将其上面和左边格子也加入正样本, 反之下面和右边
3. 根据公式计算出定位损失(长宽, 位置, iou值)
4. 根据分类, 使用BCELoss计算其分类损失
4. 将当前anchor相关层的所有预测框的iou计算出来后, 使用BCELoss计算其置信度损失
5. 相加
首先， 初始化BCELoss函数
然后，将Ground True 的真实检测框预处理成需要的格式，注意，这里真实检测框需要根据anchor的类型，判断这个样本是正样本还是负样本，所以说我们在训练时anchor的预设值是很重要的！！！
最后，分别进行定位损失，置信度损失，分类损失计算
定位损失
分类损失
置信度损失
最后相加，得到最终结果
完整代码
class ComputeLoss: # TODO 分类损失 定位损失 置信度损失
"""
计算出当前模型输出的损失值
核心逻辑 : 分类损失 + 定位损失 + 置信度损失
1. 初始化置信度损失和分类损失所需要用到的BCELoss函数
2. 得到GT(Ground True)正样本相关参数(索引, 分类, 长宽位置等等)
2.1 比较GT和anchor的宽比和高比，大于一定的阈值(4)就是负样本
2.2 筛选到的正样本信息, 观察其中心点坐标, 小于 0.5 将其上面和左边格子也加入正样本, 反之下面和右边
3. 根据公式计算出定位损失(长宽, 位置, iou值)
4. 根据分类, 使用BCELoss计算其分类损失
4. 将当前anchor相关层的所有预测框的iou计算出来后, 使用BCELoss计算其置信度损失
5. 相加
"""
sort_obj_iou = False # 后面筛选置信度损失正样本的时候是否先对iou排序
# Compute losses
def __init__(self, model, autobalance=False):
"""Initializes ComputeLoss with model and autobalance option, autobalances losses if True."""
device = next(model.parameters()).device # get model device
h = model.hyp # hyperparameters
# Define criteria
BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device)) # 分类的 BCELoss 函数
BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["obj_pw"]], device=device)) # 置信度 BCELoss 函数
# Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
# 标签平滑 eps=0代表不做标签平滑-> cp=1 cn=0 eps!=0代表做标签平滑 cp代表positive的标签值 cn代表negative的标签值
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: # g=0 代表不用focal loss
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
m = de_parallel(model).model[-1] # Detect() module
# balance用来设置三个feature map对应输出的置信度损失系数(平衡三个feature map的置信度损失)
# 从左到右分别对应大feature map(检测小目标)到小feature map(检测大目标)
# 思路: It seems that larger output layers may overfit earlier, so those numbers may need a bit of adjustment
# 一般来说，检测小物体的难度大一点，所以会增加大特征图的损失系数，让模型更加侧重小物体的检测
# 如果det.nl=3就返回[4.0, 1.0, 0.4]否则返回[4.0, 1.0, 0.25, 0.06, .02]
# self.balance = {3: [4.0, 1.0, 0.4], 4: [4.0, 1.0, 0.25, 0.06], 5: [4.0, 1.0, 0.25, 0.06, .02]}[det.nl]
self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7
# 三个预测头的下采样率det.stride: [8, 16, 32] .index(16): 求出下采样率stride=16的索引
# 这个参数会用来自动计算更新3个feature map的置信度损失系数self.balance
self.ssi = list(m.stride).index(16) if autobalance else 0 # stride 16 index
# self.BCEcls: 类别损失函数 self.BCEobj: 置信度损失函数 self.hyp: 超参数
# self.gr: 计算真实框的置信度标准的iou ratio self.autobalance: 是否自动更新各feature map的置信度损失平衡系数 默认False
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
self.na = m.na # number of anchors 3
self.nc = m.nc # number of classes 80
self.nl = m.nl # number of layers 3
self.anchors = m.anchors
self.device = device
def __call__(self, p, targets): # predictions, targets
"""Performs forward pass, calculating class, box, and object loss for given predictions and targets."""
lcls = torch.zeros(1, device=self.device) # class loss
lbox = torch.zeros(1, device=self.device) # box loss
lobj = torch.zeros(1, device=self.device) # object loss
tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets
# [[251,],[195,],[126,]] [[251,4],[195,4],[126,4]] [[[251,]...], [[195,]...],[[126,]...]] [[251,2],[195,2],[126,2]]
# Losses
for i, pi in enumerate(p): # layer index, layer predictions
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device) # target obj
n = b.shape[0] # number of targets
if n:
# pxy, pwh, _, pcls = pi[b, a, gj, gi].tensor_split((2, 4, 5), dim=1) # faster, requires torch 1.8.0
pxy, pwh, _, pcls = pi[b, a, gj, gi].split((2, 2, 1, self.nc), 1) # target-subset of predictions
# 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
# Objectness(置信度损失)
iou = iou.detach().clamp(0).type(tobj.dtype)
if self.sort_obj_iou:
j = iou.argsort()
b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
if self.gr < 1:
iou = (1.0 - self.gr) + self.gr * iou
tobj[b, a, gj, gi] = iou # iou ratio 记录iou值
# Classification 分类损失
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
# Append targets to text file
# with open('targets.txt', 'a') as file:
# [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
obji = self.BCEobj(pi[..., 4], tobj)
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 *= 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()
def build_targets(self, p, targets): # [(4,3,80,80,85),(4,3,40,40,85),(4,3,20,20,85)] (55,6)
"""Prepares model targets from input targets (image,class,x,y,w,h) for loss computation, returning class, box,
indices, and anchors.
"""
na, nt = self.na, targets.shape[0] # number of anchors, targets 3 55(4张图像所有的检测框)
tcls, tbox, indices, anch = [], [], [], []
gain = torch.ones(7, device=self.device) # normalized to gridspace gain (7,)
ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt) (3,55)
targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2) # append anchor indices (3,55,7)
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 (5,2)
# 分层进行损失计算, 筛选gt的anchor正样本
for i in range(self.nl): # 3
anchors, shape = self.anchors[i], p[i].shape # (3,2) (4,3,80,80,85)
# gain: 保存每个输出feature map的宽高 -> gain[2:6]=gain[whwh]
gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]] # xyxy gain tensor([ 1., 1., 80., 80., 80., 80., 1.])
# Match targets to anchors
t = targets * gain # shape(3,n,7) 变成在特征图上的坐标了
if nt: # 有检测框
# 所有的ground true(真实框)与当前层的三个anchor的宽高比(w/w h/h)
r = t[..., 4:6] / anchors[:, None] # wh ratio (3,55,2)
# 筛选条件 GT与anchor的宽比或高比超过一定的阈值 就当作负样本
j = torch.max(r, 1 / r).max(2)[0] < self.hyp["anchor_t"] # compare (3,55)
# j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
# 根据筛选条件j, 过滤负样本, 得到所有gt的anchor正样本(batch_size张图片)
# 知道当前gt的坐标 属于哪张图片 正样本对应的idx 也就得到了当前gt的正样本anchor
t = t[j] # filter (84,7)
# Offsets
# Offsets 筛选当前格子周围格子 找到2个离target中心最近的两个格子 可能周围的格子也预测到了高质量的样本 我们也要把这部分的预测信息加入正样本中
# 除了target所在的当前格子外, 还有2个格子对目标进行检测(计算损失) 也就是说一个目标需要3个格子去预测(计算损失)
# 首先当前格子是其中1个 再从当前格子的上下左右四个格子中选择2个 用这三个格子去预测这个目标(计算损失)
# feature map上的原点在左上角 向右为x轴正坐标 向下为y轴正坐标
gxy = t[:, 2:4] # grid xy (84,2)
gxi = gain[[2, 3]] - gxy # inverse
# 筛选中心坐标 距离当前grid_cell的左、上方偏移小于g=0.5 且 中心坐标必须大于1(坐标不能在边上 此时就没有4个格子了)
# j: [126] bool 如果是True表示当前target中心点所在的格子的左边格子也对该target进行回归(后续进行计算损失)
# k: [126] bool 如果是True表示当前target中心点所在的格子的上边格子也对该target进行回归(后续进行计算损失)
j, k = ((gxy % 1 < g) & (gxy > 1)).T
# 筛选中心坐标 距离当前grid_cell的右、下方偏移小于g=0.5 且 中心坐标必须大于1(坐标不能在边上 此时就没有4个格子了)
# l: [126] bool 如果是True表示当前target中心点所在的格子的右边格子也对该target进行回归(后续进行计算损失)
# m: [126] bool 如果是True表示当前target中心点所在的格子的下边格子也对该target进行回归(后续进行计算损失)
l, m = ((gxi % 1 < g) & (gxi > 1)).T
# j: [5, 126] torch.ones_like(j): 当前格子, 不需要筛选全是True j, k, l, m: 左上右下格子的筛选结果
j = torch.stack((torch.ones_like(j), j, k, l, m))
# 得到筛选后所有格子的正样本 格子数<=3*126 都不在边上等号成立
# t: [126, 7] -> 复制5份target[5, 126, 7] 分别对应当前格子和左上右下格子5个格子
# j: [5, 126] + t: [5, 126, 7] => t: [378, 7] 理论上是小于等于3倍的126 当且仅当没有边界的格子等号成立
t = t.repeat((5, 1, 1))[j]
# torch.zeros_like(gxy)[None]: [1, 126, 2] off[:, None]: [5, 1, 2] => [5, 126, 2]
# j筛选后: [378, 2] 得到所有筛选后的网格的中心相对于这个要预测的真实框所在网格边界（左右上下边框）的偏移量
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
扫一扫
176
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
