损失函数与标签分配
- 1 IOU损失
- 2 YOLOX的损失函数类
* - （1）初始化与forward函数
- （2）边框调整与网格生成函数self.get_output_and_grid
- （3）损失计算函数self.get_losses（上）
- （4）标签分配函数self.get_assignments（上）
- （5）第一轮筛选self.get_in_boxes_info
- （6）第二轮筛选与self.bboxes_iou、self.dynamic_k_matching（标签分配函数（下））
- （7）损失计算函数self.get_losses（下）
- （8）YOLOLoss类的测试
我们今天来介绍一下YOLOX模型的训练，我们前面已经完成了模型的推理和输出解码，现在就差损失函数了。一张图片输入模型后，会输出8400个预测框（解码只起到转化的作用，起不到过滤的作用），如何将这8400个预测框与真实目标框整合到一个损失函数里面去？YOLOv1~v5的策略是将8400个网格区域按照有没有目标划分成正负样本，然后与8400个预测框求损失，具体公式可以看代码。YOLOX摒弃这种策略，取而代之的是使用动态标签分配策略，也就是为每个标签分配若干个预测框，通过这些预测框来求损失。我们今天要介绍的就是这个标签分配策略和损失函数。
个人认为，这篇博客是”从零开始实现yolox”系列中最难，因为原作者将标签分配的很多函数都写到一个类里面了，而且像套娃一样层层调用，很难像前面的几篇文章一样，每写完一个模块就写一个测试函数去测试，只能把整个类写完后再统一写测试函数。
1 IOU损失
yolox中损失函数的公式为：
也就是说，yolox的输出中，边框的中心点和高宽使用回归损失（IOU损失），目标置信度和类别使用分类损失。这里我们定义一个名为IOUloss的类来求IOU损失。
在yolox_from_scratch/nets下新建一个名为 yolo_training.py的文件
在里面加上一个IOUloss函数
import torch
import torch.nn as nn
import torch.nn.functional as F
class IOUloss(nn.Module):
def __init__(self, reduction="none", loss_type="iou"):
"""
Args:
reduction:
loss_type:损失类型
"""
super(IOUloss, self).__init__()
self.reduction = reduction
self.loss_type = loss_type
def forward(self, pred, target):
"""
Args:
pred:预测框，维度为 (num_boxes, 4)
target:真实框，维度为 (num_boxes, 4)
Returns:
"""
assert pred.shape[0] == target.shape[0]
pred = pred.view(-1, 4)
target = target.view(-1, 4)
tl = torch.max(
(pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)
)
br = torch.min(
(pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)
)
area_p = torch.prod(pred[:, 2:], 1)
area_g = torch.prod(target[:, 2:], 1)
en = (tl < br).type(tl.type()).prod(dim=1)
area_i = torch.prod(br - tl, 1) * en
iou = (area_i) / (area_p + area_g - area_i + 1e-16)
if self.loss_type == "iou":
loss = 1 - iou ** 2
elif self.loss_type == "giou":
c_tl = torch.min(
(pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)
)
c_br = torch.max(
(pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)
)
area_c = torch.prod(c_br - c_tl, 1)
giou = iou - (area_c - area_i) / area_c.clamp(1e-16)
loss = 1 - giou.clamp(min=-1.0, max=1.0)
if self.reduction == "mean":
loss = loss.mean()
elif self.reduction == "sum":
loss = loss.sum()
return loss
代码中默认是普通的IOU损失
写一个测试脚本，yolox_from_scratch下写一个名为iou_loss_test.py的脚本
加入下面的代码：
import torch
from nets.yolo_training import IOUloss
pred = torch.tensor([150, 150, 200, 120]).reshape(-1, 4)
target = torch.tensor([300, 170, 300, 400]).reshape(-1, 4)
Iou_loss = IOUloss()
print(Iou_loss(pred, target))
输出：
tensor([0.9917])
我们可以自己画图言这个一下这个答案
2 YOLOX的损失函数类
我们定义一个名为YOLOLoss的类来求损失函数。这个类是最难的，耗费了巨大的时间，才读懂源代码。这篇文章的精华，都在函数的注释里。需要注意的是，上一篇文章，我们是使用官方预训练权重，相应的模型是对80个类别进行检测，而我们自己的数据集只有4个类别，这个区别会体现在求损失函数的过程中，各个中间量的维度里。
（1）初始化与forward函数
在 yolo_training.py中写一个名为 YOLOLoss的类及其初始化方法
class YOLOLoss(nn.Module):
def __init__(self, num_classes, strides=[8, 16, 32]):
super().__init__()
self.num_classes = num_classes
self.strides = strides
self.bcewithlog_loss = nn.BCEWithLogitsLoss(reduction="none")
self.iou_loss = IOUloss(reduction="none")
self.grids = [torch.zeros(1)] * len(strides)
下面是YOLOLoss类的forward方法：
def forward(self, inputs, labels=None):
"""
Args:
inputs:模型的输出，一个列表，列表中的元素，为各个检测头的输出，维度为(batch_size, anchor_attr, grid_w, grid_h)
各个检测头输出的张量维度，只有grid_w和grid_h不一样
labels:每次从dataloader获得的一个batch的标签，这是一个列表，每个元素对应一张图片，元素个数为batch_size，
列表中的每个元素为(num_gt, 5)，如果对应的图片中没有GT，那么labels中相应的元素就是一个空张量，否则是(num_objs, 5)
5列表示(x, y, w, h, cls_index)，cls_index是类别索引
Returns:
"""
outputs = []
x_shifts = []
y_shifts = []
expanded_strides = []
for k, (stride, output) in enumerate(zip(self.strides, inputs)):
output, grid = self.get_output_and_grid(output, k, stride)
x_shifts.append(grid[:, :, 0])
y_shifts.append(grid[:, :, 1])
expanded_strides.append(torch.ones_like(grid[:, :, 0]) * stride)
outputs.append(output)
loss = self.get_losses(x_shifts, y_shifts, expanded_strides, labels, torch.cat(outputs, 1))
return loss
（2）边框调整与网格生成函数self.get_output_and_grid
def forward中出现了两个成员函数，self.get_output_and_grid和self.get_losses，我们先来介绍第一个：
def get_output_and_grid(self, output, k, stride):
"""
在这个函数之前，第k个检测头对应的grid为tensor([0.])，预测框的中心点坐标和高宽都是相对于网格的
本函数的目的是生成一个张量来表示grid，使其表示每个网格左上角点在输出特征图中的位置
并让output的中心点坐标和高宽变成letterbox图像中的数据，这一功能和解码是一样的
下面的注释，是假设调用这个方法是(80, 80)的检测头（除了这个检测头之外，还有还有(40, 40)和(20, 20)两个检测头）
Args:
output:当前检测头的输出，维度为(batch_size, anchor_attr, grid_w, grid_h)
k:当前检测头的序号，如果只有三个检测头，那么k=0~2
stride:当前检测头输出的特征层的步长
Returns:output:预测框数据调整后的结果，已将预测框的中心点坐标和高宽转化成letterbox中的数据，
维度为(batch_size, anchor_attr, grid_w, grid_h)
grid:表示每个网格左上角点在输出特征图中的位置
"""
grid = self.grids[k]
hsize, wsize = output.shape[-2:]
if grid.shape[2:4] != output.shape[2:4]:
yv, xv = torch.meshgrid([torch.arange(hsize), torch.arange(wsize)])
grid = torch.stack((xv, yv), 2).view(1, hsize, wsize, 2).type(output.type())
self.grids[k] = grid
grid = grid.view(1, -1, 2)
output = output.flatten(start_dim=2).permute(0, 2, 1)
output[..., :2] = (output[..., :2] + grid) * stride
output[..., 2:4] = torch.exp(output[..., 2:4]) * stride
return output, grid
（3）损失计算函数self.get_losses（上）
接下来是self.get_losses，个人认为这个是最复杂的，因为从这开始调用了很多其他函数，而且是层层调用，我花了一个星期才吃透这个函数。下面这段代码，只需要先看到torch.cuda.empty_cache()
def get_losses(self, x_shifts, y_shifts, expanded_strides, labels, outputs):
"""
计算当前batch的损失，以下注释都是假设 batch_size=16，input_size=(640, 640)，num_cls=4 的情况下的注释
Args:
x_shifts:一个列表，包括了三个元素，每个元素代表了对应检测头的各个网格在x方向上的偏移，三个元素的维度分别为：(1, 6400), (1, 1600), (1, 400)
y_shifts:一个列表，包括了三个元素，每个元素代表了对应检测头的各个网格在y方向上的偏移，每个元素的维度和x_shifts一致
expanded_strides:一个包括了三个元素的列表，每个元素的维度和x_shifts对应的元素相同，表示每个网格（anchor）的步长
labels:每次从dataloader获得的一个batch的标签，这是一个列表，每个元素对应一张图片，元素个数为batch_size，
列表中的每个元素为(num_gt, 5)，如果对应的图片中没有GT，那么labels中相应的元素就是一个空张量，否则是(num_objs, 5)
5列表示(x, y, w, h, cls_index)，cls_index是类别索引
outputs:模型输出后经过torch.cat的结果，维度为 torch.Size([16, 8400, 9])，16是batch_size，8400是anchor的数量
9列表示(x, y, w, h, obj_preds, cls_preds)，cls_preds表示4个类别的概率，共4列
Returns:
"""
bbox_preds = outputs[:, :, :4]
obj_preds = outputs[:, :, 4:5]
cls_preds = outputs[:, :, 5:]
total_num_anchors = outputs.shape[1]
x_shifts = torch.cat(x_shifts, 1)
y_shifts = torch.cat(y_shifts, 1)
expanded_strides = torch.cat(expanded_strides, 1)
cls_targets = []
reg_targets = []
obj_targets = []
fg_masks = []
num_fg = 0.0
for batch_idx in range(outputs.shape[0]):
num_gt = len(labels[batch_idx])
if num_gt == 0:
cls_target = outputs.new_zeros((0, self.num_classes))
reg_target = outputs.new_zeros((0, 4))
obj_target = outputs.new_zeros((total_num_anchors, 1))
fg_mask = outputs.new_zeros(total_num_anchors).bool()
else:
gt_bboxes_per_image = labels[batch_idx][..., :4]
gt_classes = labels[batch_idx][..., 4]
bboxes_preds_per_image = bbox_preds[batch_idx]
cls_preds_per_image = cls_preds[batch_idx]
obj_preds_per_image = obj_preds[batch_idx]
gt_matched_classes, fg_mask, pred_ious_this_matching, matched_gt_inds, num_fg_img = self.get_assignments(
num_gt, total_num_anchors, gt_bboxes_per_image, gt_classes, bboxes_preds_per_image,
cls_preds_per_image, obj_preds_per_image, expanded_strides, x_shifts, y_shifts
)
torch.cuda.empty_cache()
num_fg += num_fg_img
cls_target = F.one_hot(gt_matched_classes.to(torch.int64), self.num_classes).float() * pred_ious_this_matching.unsqueeze(-1)
obj_target = fg_mask.unsqueeze(-1)
reg_target = gt_bboxes_per_image[matched_gt_inds]
cls_targets.append(cls_target)
reg_targets.append(reg_target)
obj_targets.append(obj_target.type(cls_target.type()))
fg_masks.append(fg_mask)
cls_targets = torch.cat(cls_targets, 0)
reg_targets = torch.cat(reg_targets, 0)
obj_targets = torch.cat(obj_targets, 0)
fg_masks = torch.cat(fg_masks, 0)
num_fg = max(num_fg, 1)
loss_iou = (self.iou_loss(bbox_preds.view(-1, 4)[fg_masks], reg_targets)).sum()
loss_obj = (self.bcewithlog_loss(obj_preds.view(-1, 1), obj_targets)).sum()
loss_cls = (self.bcewithlog_loss(cls_preds.view(-1, self.num_classes)[fg_masks], cls_targets)).sum()
reg_weight = 5.0
loss = reg_weight * loss_iou + loss_obj + loss_cls
return loss / num_fg
（4）标签分配函数self.get_assignments（上）
def get_losses调用了self.get_assignments，它是将8400个anchor划分成正负样本，正样本就是能和GT进行匹配的anchor，负样本就是不能和GT进行匹配的anchor，正样本可以和GT计算分类、回归、置信度损失，负样本只能计算置信度损失。这个函数的代码如下（先展示一部分，讲完第二轮筛选后会讲第二部分）：
@torch.no_grad()
def get_assignments(self, num_gt, total_num_anchors, gt_bboxes_per_image, gt_classes,bboxes_preds_per_image,
cls_preds_per_image, obj_preds_per_image, expanded_strides, x_shifts, y_shifts):
"""
一张图片输入到模型后，三个检测头会得到8400个anchor，这些anchor只有一部分会当成正样本，与标签进行损失函数计算
这个函数就是把这些anchor给找出来
Args:
num_gt:当前图片中GT的数量，纯数字
total_num_anchors:三个检测头的anchor总数，纯数字，这里是8400
gt_bboxes_per_image:当前图片中GT的中心点坐标及宽高，维度为(num_gt, 4)
gt_classes:当前图片中，所有GT的类别索引，维度为(num_gt,)
bboxes_preds_per_image:当前图片预测框的中心点坐标及宽高，维度为(8400, 4)
cls_preds_per_image:当前图片预测目标的类别，维度为(8400, 4)
obj_preds_per_image:当前图片预测目标的置信度（目标置信度），维度为(8400, 1)
expanded_strides:各个anchor与输入图片中网格的尺寸比例，即步长，维度为(1, 8400)
x_shifts:各个anchor在特征图中的横坐标，维度为(1, 8400)
y_shifts:各个anchor在特征图中的纵坐标，维度为(1, 8400)
Returns:gt_matched_classes:第二轮筛选后得到的anchor对应GT的索引，维度为(len_sg, )，len_sg是经过第二轮筛选后得到的anchor数量
fg_mask:第二轮筛选后得到的anchor在8400个anchor中的布尔索引，维度为(8400, )
pred_ious_this_matching:第二轮筛选得到的anchor，与其对应的GT的iou，维度为(len_sg, )
matched_gt_inds:第二轮筛选得到的anchor能和哪些GT匹配，维度为(len_sg, )
num_fg:当前图片中，经过两轮筛选后，所有GT的正样本总数，即能与任意一个GT匹配的anchor总数，一个纯数字
"""
"""第一轮筛选"""
fg_mask, is_in_boxes_and_center = self.get_in_boxes_info(gt_bboxes_per_image, expanded_strides, x_shifts,
y_shifts, total_num_anchors, num_gt)
（5）第一轮筛选self.get_in_boxes_info
def get_assignments出现了self.get_in_boxes_info函数，它对8400个anchor做第一轮筛选。第一轮筛选使用了两种方法，任意一个anchor只要通过其中一种筛选方法，就可以认为其通过了第一轮筛选。
第一种方法是先把网格的各个中心点坐标求出来，判断其是否在GT的内部，如果在GT的内部，那么就认为该网格对应的anchor与GT匹配。如下图所示：
第二种方法是以每个GT的中心点为中心，生成一个边长为5的正方形，判断各个网格的中心点是否在这个网格的内部
看懂了上面两幅图，就很好理解self.get_in_boxes_info函数了，现在我们来看看self.get_in_boxes_info的代码：
def get_in_boxes_info(self, gt_bboxes_per_image, expanded_strides, x_shifts, y_shifts, total_num_anchors, num_gt, center_radius = 2.5):
"""
三个检测头，共有8400个anchor，但这些anchor只有部分能和标签进行匹配，本函数就是筛选出能和标签进行匹配的anchor
本函数中使用两种方法对anchor进行筛选
Args:
gt_bboxes_per_image:当前图片中，各个真实框的中心点坐标及宽高，维度为(num_gt, 4)
expanded_strides:每个网格的步长，维度为torch.Size([1, 8400])
x_shifts:维度为torch.Size([1, 8400])
y_shifts:维度为torch.Size([1, 8400])
total_num_anchors:网格点总数，纯数字，例如8400
num_gt:真实框总数，纯数字
center_radius:半径，纯数字
Returns:is_in_boxes_anchor 能通过两种方法之一的anchor的布尔索引，维度为(8400, )
is_in_boxes_and_center 这也是一个布尔索引，表示第一种筛选方法得到的anchor中，能通过第二种筛选方法的anchor，
维度为(num_gt, len_first)，len_first是is_in_boxes_anchor中True的数量
"""
expanded_strides_per_image = expanded_strides[0]
x_centers_per_image = ((x_shifts[0] + 0.5) * expanded_strides_per_image).unsqueeze(0).repeat(num_gt, 1)
y_centers_per_image = ((y_shifts[0] + 0.5) * expanded_strides_per_image).unsqueeze(0).repeat(num_gt, 1)
"""第一种筛选方式：筛选出中心点在GT内部的网格，所对应的anchor"""
"""各个GT的上下左右边缘"""
gt_bboxes_per_image_l = (gt_bboxes_per_image[:, 0] - 0.5 * gt_bboxes_per_image[:, 2]).unsqueeze(1).repeat(1, total_num_anchors)
gt_bboxes_per_image_r = (gt_bboxes_per_image[:, 0] + 0.5 * gt_bboxes_per_image[:, 2]).unsqueeze(1).repeat(1, total_num_anchors)
gt_bboxes_per_image_t = (gt_bboxes_per_image[:, 1] - 0.5 * gt_bboxes_per_image[:, 3]).unsqueeze(1).repeat(1, total_num_anchors)
gt_bboxes_per_image_b = (gt_bboxes_per_image[:, 1] + 0.5 * gt_bboxes_per_image[:, 3]).unsqueeze(1).repeat(1, total_num_anchors)
"""计算各个网格中心点与GT各个边缘的距离"""
b_l = x_centers_per_image - gt_bboxes_per_image_l
b_r = gt_bboxes_per_image_r - x_centers_per_image
b_t = y_centers_per_image - gt_bboxes_per_image_t
b_b = gt_bboxes_per_image_b - y_centers_per_image
bbox_deltas = torch.stack([b_l, b_t, b_r, b_b], 2)
"""获得各个anchor的匹配情况"""
is_in_boxes = bbox_deltas.min(dim=-1).values > 0.0
is_in_boxes_all = is_in_boxes.sum(dim=0) > 0
"""第二种筛选方式：以GT的中心为中心，生成一个边长为5个stride的正方形（这里简称GT方框），
将中心点落在这个正方形内的网格所对应的anchor，作为与GT匹配的正样本"""
gt_bboxes_per_image_l = (gt_bboxes_per_image[:, 0]).unsqueeze(1).repeat(1, total_num_anchors) - center_radius * expanded_strides_per_image.unsqueeze(0)
gt_bboxes_per_image_r = (gt_bboxes_per_image[:, 0]).unsqueeze(1).repeat(1, total_num_anchors) + center_radius * expanded_strides_per_image.unsqueeze(0)
gt_bboxes_per_image_t = (gt_bboxes_per_image[:, 1]).unsqueeze(1).repeat(1, total_num_anchors) - center_radius * expanded_strides_per_image.unsqueeze(0)
gt_bboxes_per_image_b = (gt_bboxes_per_image[:, 1]).unsqueeze(1).repeat(1, total_num_anchors) + center_radius * expanded_strides_per_image.unsqueeze(0)
"""计算各个网格中心点与GT方框各个边缘的距离"""
c_l = x_centers_per_image - gt_bboxes_per_image_l
c_r = gt_bboxes_per_image_r - x_centers_per_image
c_t = y_centers_per_image - gt_bboxes_per_image_t
c_b = gt_bboxes_per_image_b - y_centers_per_image
center_deltas = torch.stack([c_l, c_t, c_r, c_b], 2)
"""获得各个anchor的匹配情况"""
is_in_centers = center_deltas.min(dim=-1).values > 0.0
is_in_centers_all = is_in_centers.sum(dim=0) > 0
"""将上述两种方法综合起来"""
is_in_boxes_anchor = is_in_boxes_all | is_in_centers_all
is_in_boxes_and_center = is_in_boxes[:, is_in_boxes_anchor] & is_in_centers[:, is_in_boxes_anchor]
return is_in_boxes_anchor, is_in_boxes_and_center
（6）第二轮筛选与self.bboxes_iou、self.dynamic_k_matching（标签分配函数（下））
第一轮筛选后，就要做第二轮筛选了。第二轮筛选使用简化的OTA算法，即simOTA算法，它的过程如下：
（1）计算每个anchor（经过第一轮筛选后得到的anchor）与每个GT的分类损失和iou损失，然后求和得到cost矩阵（成本函数）；
（2）在经过第一轮筛选后得到的anchor中，为每个GT找到与其有最大IOU的10个anchor，将这10个anchor对应的IOU值求和取整，即为当前GT所匹配到的anchor数量，即dynamic_k，IOU排名前dynamic_k的anchor即为和当前GT匹配的anchor。可以用如下例子理解这一过程：
（3）部分anchor可能和多个GT匹配，这是不允许的，要对这些anchor进行处理。假设经过上述操作后，某些anchor各自能和2个及2个以上的GT匹配，对每一个这样的anchor，寻找cost最小的GT作为与其匹配的GT，这里的cost是第（1）步中求得的cost。个人认为这里还是有些漏洞的，因为可能出现这种情况：经过第（2）步后，第j个anchor同时与第m和n个GT匹配，但第cost中，第j列最小的元素在第i行，i≠m且i≠n，也就是说，第i个GT的前dynamic_k个anchor中，并不包含第j个anchor。
理解了以上过程，就能更好地看懂第二轮筛选的代码了。让我们回到 def get_assignments函数中，添加以下代码，做第二轮筛选：
"""下面是第二轮筛选"""
"""获得筛选后的anchor的边框、类别概率和置信度"""
bboxes_preds_per_image = bboxes_preds_per_image[fg_mask]
cls_preds_ = cls_preds_per_image[fg_mask]
obj_preds_ = obj_preds_per_image[fg_mask]
num_in_boxes_anchor = bboxes_preds_per_image.shape[0]
"""计算IOU损失"""
pair_wise_ious = self.bboxes_iou(gt_bboxes_per_image, bboxes_preds_per_image, False)
pair_wise_ious_loss = -torch.log(pair_wise_ious + 1e-8)
"""计算分类损失"""
cls_preds_ = cls_preds_.float().unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_() * obj_preds_.unsqueeze(0).repeat(num_gt, 1, 1).sigmoid_()
gt_cls_per_image = F.one_hot(gt_classes.to(torch.int64), self.num_classes).float().unsqueeze(1).repeat(1, num_in_boxes_anchor, 1)
pair_wise_cls_loss = F.binary_cross_entropy(cls_preds_.sqrt_(), gt_cls_per_image, reduction="none").sum(-1)
del cls_preds_
cost = pair_wise_cls_loss + 3.0 * pair_wise_ious_loss + 100000.0 * (~is_in_boxes_and_center).float()
"""SimOTA求解"""
num_fg, gt_matched_classes, pred_ious_this_matching, matched_gt_inds = self.dynamic_k_matching(cost,
pair_wise_ious,
gt_classes,
num_gt, fg_mask)
del pair_wise_cls_loss, cost, pair_wise_ious, pair_wise_ious_loss
return gt_matched_classes, fg_mask, pred_ious_this_matching, matched_gt_inds, num_fg
这里出现了self.bboxes_iou和self.dynamic_k_matching函数，前者用来计算GT和预测框的IOU，后者用来求每个GT的dynamic_k。我们先讲一下self.bboxes_iou，其代码如下：
def bboxes_iou(self, bboxes_a, bboxes_b, xyxy=True):
"""
求GT与预测框（anchor）的交并比
Args:
bboxes_a: GT，维度为(num_gt, 4)
bboxes_b: 预测框，维度为(len_fg, 4)，len_fg是经过第一轮筛选后得到的anchor数量
xyxy:GT和预测框，是否为边框上下角点的坐标
Returns: iou GT和预测框的交并比，维度为(num_gt, len_fg)，例如iou[i, j]表示第i个GT和第j个预测框的交并比
"""
if bboxes_a.shape[1] != 4 or bboxes_b.shape[1] != 4:
raise IndexError
if xyxy:
tl = torch.max(bboxes_a[:, None, :2], bboxes_b[:, :2])
br = torch.min(bboxes_a[:, None, 2:], bboxes_b[:, 2:])
area_a = torch.prod(bboxes_a[:, 2:] - bboxes_a[:, :2], 1)
area_b = torch.prod(bboxes_b[:, 2:] - bboxes_b[:, :2], 1)
else:
tl = torch.max(
(bboxes_a[:, None, :2] - bboxes_a[:, None, 2:] / 2),
(bboxes_b[:, :2] - bboxes_b[:, 2:] / 2),
)
br = torch.min(
(bboxes_a[:, None, :2] + bboxes_a[:, None, 2:] / 2),
(bboxes_b[:, :2] + bboxes_b[:, 2:] / 2),
)
area_a = torch.prod(bboxes_a[:, 2:], 1)
area_b = torch.prod(bboxes_b[:, 2:], 1)
en = (tl < br).type(tl.type()).prod(dim=2)
area_i = torch.prod(br - tl, 2) * en
return area_i / (area_a[:, None] + area_b - area_i)
接下来是self.dynamic_k_matching函数的代码：
def dynamic_k_matching(self, cost, pair_wise_ious, gt_classes, num_gt, fg_mask):
"""
一个GT能和多个anchor进行匹配，但一个anchor只能和一个GT进行匹配，也就是说GT和anchor是一对多的关系
这个函数先进行第二轮筛选，获得若干个anchor，然后求这些anchor与对应GT、GT的目标类别、与所匹配GT的IOU
本函数还以传引用的方式对fg_mask进行了更新，更新后的fg_mask变成了第二轮筛选后得到的anchor在8400个anchor中的布尔索引
Args:
cost:第一轮筛选得到的anchor与GT的成本函数，维度为(num_gt, len_fg)，len_fg是8400个anchor经过第一轮筛选后得到的数量
pair_wise_ious:GT和第一轮得到的anchor的交并比，维度为(num_gt, len_fg)
gt_classes:当前图片中，所有GT的类别索引，维度为(num_gt,)
num_gt:当前图片中GT的数量，纯数字
fg_mask:第一轮筛选得到的anchor在8400个anchor中的布尔索引，维度为(8400, )
Returns: num_fg:当前图片中，所有GT的正样本总数，即能与任意一个GT匹配的anchor总数，一个纯数字
gt_matched_classes:第二轮筛选后得到的anchor对应GT的索引，维度为(len_sg, )，len_sg是经过第二轮筛选后得到的anchor数量
pred_ious_this_matching:第二轮筛选得到的anchor，与其对应的GT的iou，维度为(len_sg, )
matched_gt_inds:第二轮筛选得到的anchor能和哪些GT匹配，维度为(len_sg, )
假如len_fg=50，len_sg=20，
若 matched_gt_inds[5]=3 则表示第5个anchor（20中的第5个）匹配的GT的索引是3
gt_matched_classes[5]=2 则表示与第5个anchor（20中的第5个）匹配的GT（即索引为3的GT），其类别索引是2
pred_ious_this_matching[5]=0.53，则表示第5个anchor（20中的第5个），与其匹配的GT（即索引为3的GT）的iou为0.53
"""
"""初始化匹配矩阵"""
matching_matrix = torch.zeros_like(cost)
"""确定每个GT能匹配的anchor数量"""
n_candidate_k = min(10, pair_wise_ious.size(1))
topk_ious, _ = torch.topk(pair_wise_ious, n_candidate_k, dim=1)
dynamic_ks = torch.clamp(topk_ious.sum(1).int(), min=1)
"""给每个真实框选取k个标签进行匹配"""
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
del topk_ious, dynamic_ks, pos_idx
"""有些anchor可能同时和多个GT匹配，需要在matching_matrix中，对这些anchor进行处理"""
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.0
matching_matrix[cost_argmin, anchor_matching_gt > 1] = 1.0
"""第二轮筛选整理"""
fg_mask_inboxes = matching_matrix.sum(0) > 0.0
num_fg = fg_mask_inboxes.sum().item()
"""对fg_mask进行更新"""
fg_mask[fg_mask.clone()] = fg_mask_inboxes
"""获得第二轮筛选后得到的anchor，其所对应GT、GT的目标类别、与所匹配GT的IOU"""
matched_gt_inds = matching_matrix[:, fg_mask_inboxes].argmax(0)
gt_matched_classes = gt_classes[matched_gt_inds]
pred_ious_this_matching = (matching_matrix * pair_wise_ious).sum(0)[fg_mask_inboxes]
return num_fg, gt_matched_classes, pred_ious_this_matching, matched_gt_inds
（7）损失计算函数self.get_losses（下）
最后，我们回到def get_losses函数中，我们之前只看到 torch.cuda.empty_cache()，接下来要构建目标和损失函数了，代码如下：
num_fg += num_fg_img
cls_target = F.one_hot(gt_matched_classes.to(torch.int64), self.num_classes).float() * pred_ious_this_matching.unsqueeze(-1)
obj_target = fg_mask.unsqueeze(-1)
reg_target = gt_bboxes_per_image[matched_gt_inds]
cls_targets.append(cls_target)
reg_targets.append(reg_target)
obj_targets.append(obj_target.type(cls_target.type()))
fg_masks.append(fg_mask)
cls_targets = torch.cat(cls_targets, 0)
reg_targets = torch.cat(reg_targets, 0)
obj_targets = torch.cat(obj_targets, 0)
fg_masks = torch.cat(fg_masks, 0)
num_fg = max(num_fg, 1)
loss_iou = (self.iou_loss(bbox_preds.view(-1, 4)[fg_masks], reg_targets)).sum()
loss_obj = (self.bcewithlog_loss(obj_preds.view(-1, 1), obj_targets)).sum()
loss_cls = (self.bcewithlog_loss(cls_preds.view(-1, self.num_classes)[fg_masks], cls_targets)).sum()
reg_weight = 5.0
loss = reg_weight * loss_iou + loss_obj + loss_cls
return loss / num_fg
（8）YOLOLoss类的测试
最后，我们写一个测试脚本。在yolox_from_scratch下新建一个名为yolo_loss_test.py的脚本，建立后项目结构如下：
import random
import torch
import numpy as np
from torch.utils.data import DataLoader
from nets.yolo import YoloBody
from nets.yolo_training import YOLOLoss
from utils.dataloader import YoloDataset, yolo_dataset_collate
if __name__ == '__main__':
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
"""获得数据集类的相关初始化参数"""
train_annotation_path = '2007_train.txt'
with open(train_annotation_path) as f:
train_lines = f.readlines()
input_shape = [640, 640]
num_classes = 4
mosaic = False
mixup = False
"""建立数据集类对象"""
train_dataset = YoloDataset(train_lines, input_shape, num_classes, is_train=True, mosaic=mosaic, mixup=mixup)
"""建立导入器对象"""
gen = DataLoader(train_dataset, batch_size=16, pin_memory=True, collate_fn=yolo_dataset_collate)
"""建立模型对象"""
model = YoloBody(4, 's')
"""建立损失函数计算器"""
yolo_loss = YOLOLoss(4)
for iteration, batch in enumerate(gen):
images, targets = batch[0], batch[1]
images = torch.from_numpy(images).type(torch.FloatTensor)
targets = [torch.from_numpy(ann).type(torch.FloatTensor) for ann in targets]
outputs = model(images)
loss = yolo_loss(outputs, targets)
break
print(loss)
输出为
tensor(3596.6357, grad_fn=<DivBackward0>)
能正常输出，说明我们写的损失函数类，没有bug。
Original: https://blog.csdn.net/weixin_44457930/article/details/123066376
Author: weixin_44457930
Title: 从零开始实现yolox四：模型的训练（一）损失函数与标签分配
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/708911/
转载文章受原作者版权保护。转载请注明原作者出处！