YOLOv6-4.0部分代码阅读笔记-loss_fuseab.py
loss_fuseab.py
yolov6\models\losses\loss_fuseab.py
YOLOv6中的 loss_distill_ns 、 loss_distill 、 loss_fuseab 和 loss 的区别主要在于它们的应用场景、功能和目的。
loss_distill_ns ：
应用场景 ：主要用于小模型，如YOLOv6-N。
功能 ：结合了知识蒸馏技术，通过教师模型指导学生学习，提高小模型的性能。
目的 ：在保证速度的同时，提升小模型的检测精度。
loss_distill ：
应用场景 ：适用于所有版本的YOLOv6模型。
功能 ：同样使用知识蒸馏技术，通过教师模型指导学生模型训练，提升整体模型的性能。
目的 ：在保证速度的同时，提升模型的检测精度。
loss_fuseab ：
应用场景 ：主要用于融合特征和检测头的模型，如YOLOv6-M和YOLOv6-L。
功能 ：通过融合特征和检测头，提升模型的检测性能。
目的 ：在保证速度的同时，进一步提升模型的检测精度和性能。
loss ：
应用场景 ：适用于所有版本的YOLOv6模型。
功能 ：基本的损失函数，用于训练模型。
目的 ：确保模型的基本训练需求，是其他损失函数的基础。
目录
loss_fuseab.py
1.所需的库和模块
2.class ComputeLoss:
3.class VarifocalLoss(nn.Module):
4.class BboxLoss(nn.Module):
1.所需的库和模块
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# loss_fuseab ：
# 应用场景 ：主要用于融合特征和检测头的模型，如YOLOv6-M和YOLOv6-L。
# 功能 ：通过融合特征和检测头，提升模型的检测性能。
# 目的 ：在保证速度的同时，进一步提升模型的检测精度和性能。
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from yolov6.assigners.anchor_generator import generate_anchors
from yolov6.utils.general import dist2bbox, bbox2dist, xywh2xyxy, box_iou
from yolov6.utils.figure_iou import IOUloss
from yolov6.assigners.tal_assigner import TaskAlignedAssigner2.class ComputeLoss:
class ComputeLoss:
# 损失计算函数。
'''Loss computation func.'''
# fpn_strides ：是指特征金字塔网络（Feature Pyramid Network，FPN）中不同层级的步长（stride），它定义了不同特征层的下采样因子。在目标检测模型中，FPN用于结合不同尺度的特征图，以便捕捉到不同大小的目标。
# fpn_strides: [8, 16, 32]意味着第一个特征层（P3/8）的步长为8，它对应于原始图像下采样8倍后的特征图。第二个特征层（P4/16）的步长为16，对应于原始图像下采样16倍的特征图。第三个特征层（P5/32）的步长为32，对应于原始图像下采样32倍的特征图。
# 这些步长值对于模型在不同尺度上的特征提取和目标检测至关重要。较小的步长（如8）允许模型捕捉到更多的细节信息，适合检测较小的目标；而较大的步长（如32）则捕捉更粗糙的特征，适合检测较大的目标。通过结合这些不同尺度的特征，YOLOv6能够有效地检测各种尺寸的目标。
# 在YOLOv6的网络架构中， FPN 通过自上而下的路径和自底向上的路径（如 PANet 中的设计）来聚合特征，增强了特征的表达能力，尤其是对于小目标的检测能力。这种多尺度特征融合策略是目标检测领域中提高模型性能的关键技术之一。
# grid_cell_size ：参数指的是在目标检测过程中，每个网格单元（grid cell）所负责的图像区域的大小。这个参数对于确定边界框（bounding box）的回归目标至关重要，因为它影响着模型如何预测边界框的中心点位置。
# grid_cell_size=5.0 意味着在特征图中的每个网格单元被认为包含了5个像素的区域。在YOLO系列模型中，边界框的中心点通常是相对于其所在网格单元的中心来预测的。因此，如果一个边界框的中心点落在某个网格单元内，那么这个网格单元就负责预测这个边界框。
# grid_cell_offset ：参数通常设置为0.5，它表示网格单元（grid cell）中心点的偏移量。在目标检测模型中，尤其是像YOLO这样的单阶段检测模型，图像被划分为一个个的网格单元，每个网格单元负责预测其中心点落在该单元内的目标对象的边界框（bounding box）。
# grid_cell_offset=0.5 意味着网格单元的中心点相对于单元左上角的位置偏移了半个单元的大小。在YOLO模型中，边界框的中心点通常是相对于其所在网格单元的中心来预测的，而这个中心点的预测是基于网格单元的偏移量来确定的。
# 例如，如果一个边界框的中心点落在某个网格单元内，模型会预测这个中心点相对于该网格单元中心的偏移量（x, y坐标），然后根据这个偏移量和网格单元的步长（stride）来确定边界框在原始图像中的实际位置。
# 这个偏移量的设置有助于模型更准确地定位目标对象，尤其是在目标对象的尺寸接近网格单元大小时。通过这种方式，模型可以更灵活地适应不同尺寸的目标对象，提高检测的准确性。在YOLOv6的实现中，这种偏移量的设置是网格单元机制的一部分，它允许模型在保持计算效率的同时，对目标对象进行更精细的定位。
# num_classes=80 ：这个参数指的是模型需要识别的目标类别的数量。
# ori_img_size=640 ：表示模型在训练和推理时使用的原始图像尺寸。
# warmup_epoch ：参数定义了在训练开始时使用 warm-up 策略的轮数。如果设置为0，意味着从训练的第一周期（epoch）开始就使用配置文件中指定的完整学习率，而不进行任何预热。
# se_dfl=True ：表示启用分布式焦点损失（Distributed Focal Loss，简称DFL）。 DFL 是一种改进的损失函数，旨在解决目标检测任务中类别不平衡的问题，特别是在前景和背景类别之间存在显著不平衡时。
# DFL的核心思想是将目标框的位置预测视为一个概率分布，而不是单一的确定性值。这种方法可以提高模型对目标框位置的预测精度，尤其是在目标框与真实框之间存在较大不确定性时。
# reg_max = 16 ：指的是在训练过程中，每个网格单元（grid cell）最多可以预测16个边界框（bounding boxes）。这意味着每个网格单元最多可以生成16个候选框，用于检测图像中的目标对象。
# 这个参数的设置有助于控制模型的复杂度和计算量，确保模型在训练过程中不会过于复杂，同时也能在一定程度上提高检测的准确性。
def __init__(self,
fpn_strides=[8, 16, 32],
grid_cell_size=5.0,
grid_cell_offset=0.5,
num_classes=80,
ori_img_size=640,
warmup_epoch=0,
use_dfl=True,
reg_max=16,
iou_type='giou',
loss_weight={
'class': 1.0,
'iou': 2.5,
'dfl': 0.5}
):
self.fpn_strides = fpn_strides
self.grid_cell_size = grid_cell_size
self.grid_cell_offset = grid_cell_offset
self.num_classes = num_classes
self.ori_img_size = ori_img_size
self.warmup_epoch = warmup_epoch
# class TaskAlignedAssigner(nn.Module):
# -> def __init__(self, topk=13, num_classes=80, alpha=1.0, beta=6.0, eps=1e-9):
# -> def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt):
# -> 指定目标类别、指定目标边界框、经过标准化的指定目标分数、总掩膜位置(bool值)。
# -> return target_labels, target_bboxes, target_scores, fg_mask.bool()
self.formal_assigner = TaskAlignedAssigner(topk=26, num_classes=self.num_classes, alpha=1.0, beta=6.0)
self.use_dfl = use_dfl # se_dfl=True 表示启用分布式焦点损失（Distributed Focal Loss，简称DFL）。
self.reg_max = reg_max
self.proj = nn.Parameter(torch.linspace(0, self.reg_max, self.reg_max + 1), requires_grad=False)
self.iou_type = iou_type
# class VarifocalLoss(nn.Module):
# -> def forward(self, pred_score,gt_score, label, alpha=0.75, gamma=2.0):
# -> loss = (F.binary_cross_entropy(pred_score.float(), gt_score.float(), reduction='none') * weight).sum()
# -> return loss
self.varifocal_loss = VarifocalLoss().cuda()
# class BboxLoss(nn.Module):
# -> def __init__(self, num_classes, reg_max, use_dfl=False, iou_type='giou'):
# -> def forward(self, pred_dist, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask):
# -> return loss_iou, loss_dfl
self.bbox_loss = BboxLoss(self.num_classes, self.reg_max, self.use_dfl, self.iou_type).cuda()
self.loss_weight = loss_weight
def __call__(
self,
outputs,
targets,
epoch_num,
step_num
):
# feats ：特征图（Feature Maps）的列表。在卷积神经网络（CNN）中，特征图是中间层的输出，它们捕获了输入图像在不同尺度和层次上的特征。在目标检测模型中，这些特征图通常被用于不同尺度的目标检测。
# pred_scores ：预测的类别分数。这些分数表示模型对于每个预测边界框属于各个类别的置信度。 pred_scores 的形状通常是 (batch_size, num_anchors, num_classes) ，其中 batch_size 是批次大小， num_anchors 是每个特征图单元格上的锚点数量， num_classes 是类别的数量。
# pred_distri ：预测的距离分布。这是YOLOv6中DFL（Distribution Focal Loss）使用的输出，它表示模型对于每个预测边界框的位置的预测。
# pred_distri 的形状可能是 (batch_size, num_anchors, 4 * (self.reg_max + 1)) ，其中 4 表示边界框的四个坐标（左、上、右、下）， self.reg_max + 1 是每个坐标上的预测点数。
# 在YOLOv6中，这些输出通常被用于以下步骤：
# 使用 feats 进行边界框的预测。
# 使用 pred_scores 计算分类损失，并进行类别预测。
# 使用 pred_distri 计算边界框回归损失，并进行边界框的精确定位。
feats, pred_scores, pred_distri = outputs
# def generate_anchors(feats, fpn_strides, grid_cell_size=5.0, grid_cell_offset=0.5, device='cpu', is_eval=False, mode='af'):
# -> 根据特征生成锚点。
# -> return anchors, anchor_points, num_anchors_list, stride_tensor
# 在目标检测模型中，生成锚点（anchors）是将模型的输出转换为边界框预测的关键步骤。 generate_anchors 函数通常负责生成这些锚点，以及相关的辅助信息。
# 以下是这个函数调用中各个参数的可能含义：
# 1. feats ：特征图列表。这些特征图是模型的输出，通常来自不同尺度的卷积层，用于检测不同大小的目标。
# 2. self.fpn_strides ：特征金字塔网络（Feature Pyramid Network, FPN）的步幅列表。步幅定义了每个特征图单元格在原始图像中的大小。
# 3. self.grid_cell_size ：网格单元的大小。这个参数定义了在特征图上每个单元格对应于原始图像的尺寸。
# 4. self.grid_cell_offset ：网格单元的偏移量。这个参数用于调整锚点在特征图上的位置，通常用于确保锚点的中心与目标的中心对齐。
# 5. device ：用于生成锚点的设备，通常是CPU或GPU。
# 6. is_eval ：一个布尔值，指示当前是否处于评估（evaluation）模式。在评估模式下，某些操作可能会有所不同。
# 7. mode ：模式参数，可能用于指定生成锚点的具体方式或策略，例如'ab'可能代表某种特定的锚点生成模式。
# 函数的返回值包括：
# 1. anchors ：生成的锚点列表。每个锚点是一个边界框，通常由四个值组成：(x1, y1, x2, y2)，表示边界框的左上角和右下角坐标。
# 2. anchor_points ：锚点的中心点坐标。这些点通常用于定义边界框的中心位置。
# 3. n_anchors_list ：每个特征图上的锚点数量列表。这个列表可能用于后续的损失计算或其他操作。
# 4. stride_tensor ：步幅张量，包含了每个特征图单元格在原始图像中的尺寸信息。
# 在YOLOv6中，锚点的生成对于模型的性能至关重要，因为它们直接影响到目标检测的准确性和效率。通过精心设计的锚点，模型可以更有效地预测不同大小和形状的目标物体的位置。
anchors, anchor_points, n_anchors_list, stride_tensor = \
generate_anchors(feats, self.fpn_strides, self.grid_cell_size, self.grid_cell_offset, device=feats[0].device, is_eval=False, mode='ab')
assert pred_scores.type() == pred_distri.type()
# gt_bboxes_scale 是一个张量，它被初始化为与预测类别分数 pred_scores 相同的数据类型，并且填充了原始图像的大小。这个张量通常用于后续的损失计算，特别是在处理边界框回归时。
# torch.full((1,4), self.ori_img_size) ：这个操作创建了一个形状为 (1, 4) 的张量，其中每个元素都被填充为 self.ori_img_size 。
# 这段代码的结果是 gt_bboxes_scale 张量，它的形状是 (1, 4) ，每个元素都是 self.ori_img_size 。这个张量可能用于规范化或缩放边界框的坐标，或者在损失函数中作为参考值。
# 例如，如果 self.ori_img_size 是 640（图像的宽度或高度），那么 gt_bboxes_scale 将是：
# gt_bboxes_scale = torch.tensor([[640, 640, 640, 640]])
gt_bboxes_scale = torch.full((1,4), self.ori_img_size).type_as(pred_scores)
batch_size = pred_scores.shape[0]
# targets 目标。
# def preprocess(self, targets, batch_size, scale_tensor): -> return targets
targets =self.preprocess(targets, batch_size, gt_bboxes_scale)
gt_labels = targets[:, :, :1]
gt_bboxes = targets[:, :, 1:] #xyxy
# 这行代码的作用是创建一个掩码（mask），用于识别哪些真实目标框（ground truth bounding boxes）是有效的。
# gt_bboxes.sum(-1, keepdim=True) ：这个操作是对 gt_bboxes 张量沿着最后一个维度（即每个边界框的四个坐标值）进行求和。 keepdim=True 参数确保结果保持原有的维度，这样我们得到的是一个形状为 [batch_size, num_gt, 1] 的张量。
# 在YOLOv6中，这个掩码 mask_gt 用于在训练过程中忽略那些无效的目标框。例如，如果一个目标框的宽度或高度为零，或者如果它完全在图片外面，那么它就不会对模型的训练做出贡献。
# 通过使用这个掩码，我们可以确保只有有效的目标框被用于计算损失函数，从而提高模型的训练效率和准确性。
# 总结来说，这行代码的作用是生成一个掩码张量 mask_gt ，用于指示每个目标的边界框是否有效，这对于目标检测模型的训练至关重要。
mask_gt = (gt_bboxes.sum(-1, keepdim=True) > 0).float()
# pboxes pboxes 通常指的是预测框（prediction boxes），它们是由模型在特征图上生成的，用于表示检测到的目标的位置和大小。
# 在YOLOv6或任何基于锚框（anchor boxes）的目标检测模型中， anchor_points 通常指的是在特征图（feature map）上用于初始化边界框预测的点。这些点是锚框的中心点，用于确定边界框的位置。 anchor_points_s 是调整后的锚点中心点，它们被调整以适应不同尺度的特征图。
# anchor_points ：这是一个张量，包含了特征图上每个锚框的中心点坐标。在YOLO模型中，特征图的每个单元格（cell）会预测多个边界框，而 anchor_points 就是这些单元格的中心点坐标。
# tride_tensor ：这是一个张量，包含了特征图相对于原图的步长（stride）。步长是原图到特征图的下采样比例，它决定了特征图上每个单元格代表原图的区域大小。步长对于确定边界框在原图中的准确位置至关重要。
# anchor_points / stride_tensor ：这个操作是将每个锚点的坐标除以对应的步长。这样做是为了将特征图上的坐标转换回原图的坐标。因为特征图是原图的下采样版本，所以需要通过步长来调整坐标，以确保预测的边界框在原图中的位置是正确的。
# anchor_points_s ：这是调整后的锚点中心点张量，它包含了转换后的坐标。这些坐标用于与目标框（ground truth boxes）的中心点进行比较，以计算损失函数，或者用于在特征图上定位预测的边界框。
# 总结来说，这行代码的作用是将特征图上的锚点中心点坐标调整为原图坐标，这是目标检测模型中一个重要的步骤，因为它确保了模型能够准确预测边界框在原图中的位置。这种坐标转换是必要的，因为模型在特征图上进行预测，但最终的目标是在原图上定位目标。
anchor_points_s = anchor_points / stride_tensor
# 这行代码是将预测的边界框的中心点坐标（ pred_distri[..., :2] ）调整为相对于原图的坐标。这里， pred_distri 表示模型预测的边界框信息，而 anchor_points_s 是特征图上每个单元格的中心点坐标，已经通过步长调整为原图的坐标。
# pred_distri[..., :2] ：这个操作选择了 pred_distri 张量中每个边界框的前两个坐标值，即中心点的 x 和 y 坐标。这些坐标值通常是相对于特征图的坐标。
# 这行代码的作用是将模型预测的边界框的中心点坐标调整为原图坐标，这是目标检测模型中一个重要的步骤，因为它确保了模型能够准确预测边界框在原图中的位置。
# 这种坐标转换是必要的，因为模型在特征图上进行预测，但最终的目标是在原图上定位目标。通过这种方式，YOLOv6模型能够更准确地定位和识别图像中的目标。
pred_distri[..., :2] += anchor_points_s
pred_bboxes = xywh2xyxy(pred_distri)
try:
# class TaskAlignedAssigner(nn.Module):
# -> def __init__(self, topk=13, num_classes=80, alpha=1.0, beta=6.0, eps=1e-9):
# -> def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt):
# -> 指定目标类别、指定目标边界框、经过标准化的指定目标分数、总掩膜位置(bool值)。
# -> return target_labels, target_bboxes, target_scores, fg_mask.bool()
# 用于执行锚点分配（anchor assignment），这是训练过程中的一个关键步骤。
# 1. self.formal_assigner ： 这是 YOLOv6 模型中的一个方法，用于正式地将预测的锚点（anchors）分配给真实边界框（ground truth boxes）。这个方法通常包含了匹配算法，用于确定哪些预测的锚点应该与哪些真实边界框相匹配。
# 2. pred_scores.detach() ： pred_scores 是模型预测的得分，表示模型对于每个预测边界框包含目标的置信度。
# .detach() 方法用于从当前计算图中分离出这些预测得分，这意味着在反向传播时不会计算这些得分的梯度。这通常用于防止梯度回传到模型的某些部分，特别是在处理不需要更新的中间结果时。
# 3. pred_bboxes.detach() * stride_tensor ： pred_bboxes 是模型预测的边界框坐标。 .detach() 同样用于阻止梯度回传。 stride_tensor 是一个张量，用于调整预测边界框的尺度，使其与原始图像的尺度相匹配。这通常是因为在特征图上进行的预测需要转换回原始图像空间。
# 4. anchor_points ： 这是一组预定义的点，用于生成锚点（anchor boxes）。这些点定义了在特征图上每个位置的锚点的中心。
# 5. gt_labels ： 真实边界框对应的标签，用于分类任务。
# 6. gt_bboxes ： 真实边界框的坐标，用于回归任务。
# 7. mask_gt ： 掩码 mask_gt 用于在训练过程中忽略那些无效的目标框。例如，如果一个目标框的宽度或高度为零，或者如果它完全在图片外面，那么它就不会对模型的训练做出贡献。
# 输出变量：
# target_labels ：每个预测框的目标标签。
# target_bboxes ：每个预测框的目标边界框，用于计算回归损失。
# target_scores ：每个预测框的目标得分，用于计算分类损失。
# fg_mask ：前景掩码，指示哪些预测框应该被用于损失函数的计算，即那些与真实框匹配的预测框。
target_labels, target_bboxes, target_scores, fg_mask = \
self.formal_assigner(
pred_scores.detach(),
pred_bboxes.detach() * stride_tensor,
anchor_points,
gt_labels,
gt_bboxes,
mask_gt)
except RuntimeError:
print(
"OOM RuntimeError is raised due to the huge memory cost during label assignment. \
CPU mode is applied in this batch. If you want to avoid this issue, \
try to reduce the batch size or image size."
)
torch.cuda.empty_cache()
print("------------CPU Mode for This Batch-------------")
_pred_scores = pred_scores.detach().cpu().float()
_pred_bboxes = pred_bboxes.detach().cpu().float()
_anchor_points = anchor_points.cpu().float()
_gt_labels = gt_labels.cpu().float()
_gt_bboxes = gt_bboxes.cpu().float()
_mask_gt = mask_gt.cpu().float()
_stride_tensor = stride_tensor.cpu().float()
target_labels, target_bboxes, target_scores, fg_mask = \
self.formal_assigner(
_pred_scores,
_pred_bboxes * _stride_tensor,
_anchor_points,
_gt_labels,
_gt_bboxes,
_mask_gt)
target_labels = target_labels.cuda()
target_bboxes = target_bboxes.cuda()
target_scores = target_scores.cuda()
fg_mask = fg_mask.cuda()
#Dynamic release GPU memory 动态释放GPU内存
if step_num % 10 == 0:
# torch.cuda.empty_cache()
# 释放 CUDA 缓存分配器中当前持有的所有未占用的缓存内存。这使得这些内存可以在其他 GPU 应用程序中使用，并且在 nvidia-smi 命令中可见。
torch.cuda.empty_cache()
# rescale bbox 重新调整bbox
# 在目标检测框架中，如 YOLO 或 Faster R-CNN，预测的边界框（bounding boxes）通常是相对于特征图（feature map）的坐标。特征图是输入图像经过一系列卷积层和池化层后的结果，其空间分辨率通常低于原始输入图像。
# 因此，预测的边界框坐标需要被转换回原始图像的坐标空间。 target_bboxes /= stride_tensor 这行代码的作用就是将预测的边界框坐标除以 stride_tensor ，以完成这种转换。这里的 stride_tensor 表示特征图相对于原始图像的下采样（downsampling）因子。
target_bboxes /= stride_tensor
# cls loss 类别预测损失
# 这行代码的作用是：对于 fg_mask 中标记为前景的锚点，保留它们的原始标签（ target_labels 中的值）；对于 fg_mask 中标记为背景的锚点，将它们的标签替换为背景类的标签（ self.num_classes ）。
# 这种操作在目标检测中很常见，用于确保背景锚点的标签与模型的输出一致，从而在训练过程中正确地计算损失。通过这种方式，模型可以区分前景和背景，从而更准确地进行目标检测。
target_labels = torch.where(fg_mask > 0, target_labels, torch.full_like(target_labels, self.num_classes))
# torch.nn.functional.one_hot(tensor: torch.Tensor, num_classes: int = -1) -> torch.Tensor
# 在 PyTorch 中， F.one_hot 函数是 torch.nn.functional 模块下的一个函数，用于将整数类型的索引张量转换为独热编码（one-hot encoding）形式的张量。
# 这种编码方式在处理分类问题时非常有用，因为它可以将类别标签转换为一种形式，使得每个类别由一个唯一的二进制向量表示，向量中只有一个位置是1，其余位置都是0。
# 参数
# tensor：要编码的整数类型张量（Tensor）。这些整数表示类别的索引。
# num_classes：整数，指定独热编码的类别数。如果设置为-1（默认值），则类别数将被设置为输入张量中最大值加1。
# 返回值
# 返回一个形状为 (tensor.shape, num_classes) 的张量，其中每个位置的值都是0，除了最后一个维度的索引与输入张量的相应值匹配的位置，这些位置的值为1。
# 在 PyTorch 中， F.one_hot 函数用于将整数标签转换为独热编码（one-hot encoding）的形式。独热编码是一种将分类变量转换为机器学习算法可以更好处理的形式的方法。在目标检测和分类任务中，这种编码方式常用于表示类别标签。
# target_labels.long() ： target_labels 是一个包含类别标签的张量。 .long() 方法将 target_labels 转换为长整型（ torch.long ），这是因为 F.one_hot 需要整数类型的输入。
# F.one_hot(target_labels.long(), self.num_classes + 1) ： F.one_hot 是 PyTorch 的函数，它接受两个参数：整数标签和编码的维度数。 self.num_classes + 1 表示编码的维度数。通常， self.num_classes 是数据集中的类别数，而 +1 是为了包括背景或“无目标”类别。
# [..., :-1] ： 这是 PyTorch 的高级索引，用于选择除了最后一个元素之外的所有元素。 :-1 表示不包括最后一个维度，因此这里移除了为背景类别创建的额外维度。
# 整体来看，这行代码的作用是将 target_labels 中的每个类别标签转换为独热编码形式，但不包括表示背景的最后一个类别。这可能是因为在目标检测中，背景类别通常不参与分类损失的计算。
one_hot_label = F.one_hot(target_labels.long(), self.num_classes + 1)[..., :-1]
# class VarifocalLoss(nn.Module):
# -> def forward(self, pred_score,gt_score, label, alpha=0.75, gamma=2.0):
# -> loss = (F.binary_cross_entropy(pred_score.float(), gt_score.float(), reduction='none') * weight).sum()
# -> return loss
loss_cls = self.varifocal_loss(pred_scores, target_scores, one_hot_label)
target_scores_sum = target_scores.sum()
# avoid devide zero error, devide by zero will cause loss to be inf or nan.
# if target_scores_sum is 0, loss_cls equals to 0 alson
if target_scores_sum > 0:
loss_cls /= target_scores_sum
# bbox loss 边界框损失。
# class BboxLoss(nn.Module):
# -> def __init__(self, num_classes, reg_max, use_dfl=False, iou_type='giou'):
# -> def forward(self, pred_dist, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask):
# -> return loss_iou, loss_dfl
loss_iou, loss_dfl = self.bbox_loss(pred_distri, pred_bboxes, anchor_points_s, target_bboxes,
target_scores, target_scores_sum, fg_mask)
loss = self.loss_weight['class'] * loss_cls + \
self.loss_weight['iou'] * loss_iou + \
self.loss_weight['dfl'] * loss_dfl
return loss, \
torch.cat(((self.loss_weight['iou'] * loss_iou).unsqueeze(0),
(self.loss_weight['dfl'] * loss_dfl).unsqueeze(0),
(self.loss_weight['class'] * loss_cls).unsqueeze(0))).detach()
def preprocess(self, targets, batch_size, scale_tensor):
targets_list = np.zeros((batch_size, 1, 5)).tolist()
# 这段代码是一个 Python 循环，它处理了一个名为 targets 的 PyTorch 张量。这个张量被假定包含了多个目标的信息，其中每个目标的信息是一个元组或列表。代码的目的是将这些目标根据它们所属的类别索引分配到不同的列表中。
# targets.cpu().numpy().tolist() ：这个链式调用执行了以下步骤： .cpu() ：如果 targets 张量在 GPU 上，这个操作将其移动到 CPU。 .numpy() ：将 PyTorch 张量转换为 NumPy 数组。 .tolist() ：将 NumPy 数组转换为 Python 列表。
# enumerate(targets) ： enumerate 函数用于遍历 targets 列表，并且为每个元素提供一个计数器 i （从 0 开始）和该元素本身 item 。
# int(item[0]) ：这获取了每个目标信息元组或列表的第一个元素，并将其转换为整数。这个整数代表了目标所属的类别索引。
# targets_list[int(item[0])].append(item[1:]) ：这将目标信息（除了类别索引之外的部分）添加到 targets_list 中对应的列表里。 item[1:] 表示除了类别索引之外的目标信息（例如，边界框坐标、置信度等）。
# targets_list 是一个列表的列表，其中外层列表的每个元素都是一个列表，用于存储属于特定类别的目标信息。例如，如果 targets_list 的第 0 个元素是包含所有属于类别 0 的目标信息的列表，那么 targets_list[1] 将是包含所有属于类别 1 的目标信息的列表，以此类推。
# 这段代码的执行结果是一个按类别索引组织的目标信息列表，可以用于后续的处理，例如计算损失函数或进行非极大值抑制（NMS）。这种按类别组织数据的方式在目标检测任务中很常见，因为它允许模型分别处理每个类别的目标。
for i, item in enumerate(targets.cpu().numpy().tolist()):
targets_list[int(item[0])].append(item[1:])
# 这行代码用于计算 targets_list 中所有子列表的长度，并找出最长的那个长度，将其赋值给变量 max_len 。
# (len(l) for l in targets_list) ：这是一个生成器表达式，它遍历 targets_list 中的每个子列表 l ，并为每个子列表产生一个长度值。
# max((len(l) for l in targets_list)) ： max 函数接受一个可迭代对象（在这个例子中是由生成器表达式产生的值），并返回其中的最大值。在这个上下文中， max 函数将返回 targets_list 中所有子列表的最大长度。
# 这个长度信息可以用于后续的处理，例如在创建损失函数的标签时，可能需要将所有子列表填充到相同的长度，以便能够使用批处理。在这种情况下， max_len 将指示需要填充到的目标长度。
max_len = max((len(l) for l in targets_list))
# np.array(list(map(lambda l:l + [[-1,0,0,0,0]]*(max_len - len(l)), targets_list))) ：这部分代码首先通过 map 函数对 targets_list 中的每个元素应用一个匿名函数（lambda）。
# 这个匿名函数将每个元素（假设是一个列表）与 [[-1,0,0,0,0]] 列表重复 max_len - len(l) 次（即如果列表长度小于 max_len ，则用 [[-1,0,0,0,0]] 补齐到 max_len ）的结果合并。然后将这些结果转换为一个NumPy数组。
# torch.from_numpy(...) ：将NumPy数组转换为PyTorch张量（tensor）。
# [:,1:,:] ：这个切片操作丢弃了张量的第一列（索引为0的列）。
# .to(targets.device) ：将张量移动到 targets 张量所在的设备（例如，如果 targets 在GPU上，这个操作会将新张量也移动到GPU）。
# 总的来说，这段代码的目的是将一个列表的列表（ targets_list ），每个内部列表代表一个序列，转换为一个固定长度的序列张量，不足的部分用 [[-1,0,0,0,0]] 补齐，然后丢弃第一列，并将结果移动到指定的设备上。
# 这在处理序列数据时很常见，需要将不同长度的序列转换为固定长度以适应模型输入。
targets = torch.from_numpy(np.array(list(map(lambda l:l + [[-1,0,0,0,0]]*(max_len - len(l)), targets_list)))[:,1:,:]).to(targets.device)
# batch_target 变量的这一行代码执行的操作是将目标张量（ targets ）中的特定部分与一个缩放张量（ scale_tensor ）相乘，并且是原地操作（in-place），意味着直接修改 targets 张量。
# targets[:, :, 1:5] ：这个操作选择了 targets 张量中每个目标的第1到第4个元素（不包括索引5）。在YOLO模型中，这些元素通常代表目标框的中心点坐标（x, y）和目标框的宽度和高度（w, h），它们都是相对于图片宽度和高度的比例值。
# mul_(scale_tensor) ： scale_tensor 是一个缩放因子张量，它的作用是将上述目标框的坐标和尺寸按比例缩放，以适应不同的图片尺寸或者网络输出层的分辨率。这个缩放操作是原地进行的，意味着 targets 张量中对应的元素会被更新。
# 在YOLOv6中，这种缩放操作是必要的，因为模型可能会在不同尺寸的图片上进行训练或推理，而目标框的坐标和尺寸需要根据实际图片的尺寸进行调整。通过这种方式，模型可以保持对不同尺寸图片的鲁棒性。
# 总结来说，这行代码的作用是将 targets 张量中每个目标的中心点坐标和尺寸乘以一个缩放因子，以确保目标框的坐标和尺寸与实际图片尺寸相匹配。
batch_target = targets[:, :, 1:5].mul_(scale_tensor)
# def xywh2xyxy(bboxes): -> 将 bbox(xywh) 转换为 box(xyxy)。 -> return bboxes
targets[..., 1:] = xywh2xyxy(batch_target)
return targets
def bbox_decode(self, anchor_points, pred_dist):
if self.use_dfl:
batch_size, n_anchors, _ = pred_dist.shape
pred_dist = F.softmax(pred_dist.view(batch_size, n_anchors, 4, self.reg_max + 1), dim=-1).matmul(self.proj.to(pred_dist.device))
return dist2bbox(pred_dist, anchor_points)3.class VarifocalLoss(nn.Module):
class VarifocalLoss(nn.Module):
def __init__(self):
super(VarifocalLoss, self).__init__()
def forward(self, pred_score,gt_score, label, alpha=0.75, gamma=2.0):
weight = alpha * pred_score.pow(gamma) * (1 - label) + gt_score * label
# torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)
# torch.cuda.amp.autocast 是 PyTorch 中用于自动混合精度（Automatic Mixed Precision, AMP）训练的上下文管理器。它允许模型在训练过程中动态选择使用单精度（FP32）或半精度（FP16）进行计算，以此来平衡计算速度、内存使用和数值精度。
# 参数说明
# enabled (bool): 是否启用自动混合精度。默认为 True 。
# dtype (torch.dtype): 指定自动转换的目标数据类型。如果指定，就以该类型为准；如果没有指定，则根据 device_type 为 "cuda" 时默认为 torch.float16 ，为 "cpu" 时默认为 torch.bfloat16 。
# cache_enabled (bool): 是否启用缓存。默认为 True 。当启用时，可以提高性能，但可能会增加一些内存开销。
# 工作原理
# 上下文管理器： autocast 可以作为上下文管理器使用，它会影响代码块内的运算，使其在 FP16 或 FP32 之间自动切换。 with autocast():
# 装饰器： autocast 也可以作为装饰器使用，应用于模型的 forward 方法或其他需要自动混合精度的函数。 @autocast()
with torch.cuda.amp.autocast(enabled=False):
# torch.nn.functional.binary_cross_entropy(input, target, weight=None, size_average=None, reduce=None, reduction='mean')
# 参数说明
# input (torch.Tensor) ：模型的原始输出，即logits，未经过 sigmoid 函数处理的值。这个张量的形状通常是[batch_size, ...]，其中...可以是任何数量的额外维度。
# target (torch.Tensor) ：真实标签，形状与 input 相同。标签的值应该是0或1，分别表示负类和正类。
# weight (torch.Tensor,可选) ：每个样本的权重。如果提供，它的形状应该是[batch_size, ...]，与 input 相同。这可以用来处理不平衡的数据集，给某些类别的样本更高的权重。
# size_average (bool, 可选) ：这个参数已经被弃用。如果你想要控制损失的约简方式，应该使用 reduction 参数。
# reduce (bool, 可选) ：这个参数已经被弃用。如果你想要控制损失的约简方式，应该使用 reduction 参数。
# reduction (str, 可选) ：指定如何约简损失。选项包括 ：
# 'mean' ：输出所有损失的平均值（默认）。
# 'sum' ：输出所有损失的总和。
# 'none' ：不进行任何约简，输出每个样本的损失。
# 返回值
# torch.Tensor ：计算得到的损失值。
# 如果 reduction 是 'mean' ，则返回一个标量；如果 reduction 是 'sum' ，则返回一个与 input 相同形状的张量；如果 reduction 是 'none' ，则返回一个与 input 相同形状的张量，其中每个元素是对应样本的损失。
loss = (F.binary_cross_entropy(pred_score.float(), gt_score.float(), reduction='none') * weight).sum()
return loss
4.class BboxLoss(nn.Module):
class BboxLoss(nn.Module):
def __init__(self, num_classes, reg_max, use_dfl=False, iou_type='giou'):
super(BboxLoss, self).__init__()
self.num_classes = num_classes
# class IOUloss: -> 计算IoU损失。 -> def __call__(self, box1, box2): -> return loss
self.iou_loss = IOUloss(box_format='xyxy', iou_type=iou_type, eps=1e-10)
self.reg_max = reg_max
self.use_dfl = use_dfl
# pred_dist ：是指预测边界框（bounding box）与目标物体实际边界框之间的距离。
# pred_dist 通常指的是预测的边界框的四个距离值：左（left）、上（top）、右（right）、下（bottom）。这些距离是相对于特征图上的锚点（anchor point）的，锚点通常是特征图上每个单元格的中心。模型会预测这些距离，然后通过解码过程将这些距离转换为边界框的坐标（x1, y1, x2, y2）。
# pred_bboxes ：指的是模型预测的边界框（bounding boxes）。这些预测框是模型对于输入图像中目标物体位置的估计，它们通常以相对于锚点（anchor point）的距离形式存在。
# anchor_points ：是指在特征图上用于生成预测边界框的参考点。这些点通常位于特征图的每个单元格的中心，用于作为边界框预测的起始点。在YOLOv6中，由于采用了anchor-free的设计，这些点不再与特定的锚框（anchor boxes）相关联，而是直接用于生成预测框。
# target_bboxes ：指的是在训练过程中，用于计算损失函数的真实边界框（Ground Truth Bounding Boxes）。这些边界框是从标注数据中获取的，表示了图像中每个目标物体的实际位置。
# target_scores ：是指在训练过程中，每个预测框对应的目标类别的分数。这些分数通常是基于模型对每个预测框包含特定类别物体的置信度的预测。在训练过程中，这些分数会被用来计算分类损失，以优化模型对目标类别的识别能力。
# target_scores_sum ：是指所有 target_scores 的总和。在计算分类损失时，这个总和被用作一个规范化因子，以确保损失值不会因为批次中目标框的数量不同而产生过大的波动。通过将分类损失除以 target_scores_sum ，可以得到一个平均损失值，这有助于模型在不同大小的批次上保持稳定的训练动态。
# fg_mask （foreground mask） ：是一个重要的组件，用于在训练过程中区分前景（即包含目标物体的区域）和背景。具体来说， fg_mask 是一个二进制掩码（binary mask），
# 其作用是标识出哪些预测框（anchors或prior boxes）与真实目标物体（ground truth boxes）有足够高的交并比（IoU），因此应该被视为正样本（positive samples）进行训练。
# 在目标检测的训练过程中，模型需要同时学习如何正确分类目标物体以及如何准确地定位这些物体。 fg_mask 就是用来帮助模型区分哪些预测框是有可能包含目标物体的，即前景，以及哪些预测框不包含目标物体，即背景。这样，模型就可以专注于优化那些更有可能包含目标物体的预测框。
def forward(self, pred_dist, pred_bboxes, anchor_points,
target_bboxes, target_scores, target_scores_sum, fg_mask):
# 选择正样本掩码。
# select positive samples mask
num_pos = fg_mask.sum()
if num_pos > 0:
# iou loss Iou损失。
# fg_mask 是一个二维张量，形状是 (batch_size, num_anchors) ，其中 num_anchors 是每个特征图单元格上的锚点（anchor）数量。 unsqueeze(-1) 操作在张量的最后一个维度上增加一个维度，使其形状变为 (batch_size, num_anchors, 1) 。
# repeat([1, 1, 4]) ： repeat 操作将张量中的每个元素沿着指定的维度重复指定的次数。在这里， [1, 1, 4] 表示在第一个和第二个维度（即 batch_size 和 num_anchors ）上不重复，在最后一个维度（即新增加的维度）上重复4次。
# 这样， bbox_mask 的形状就变成了 (batch_size, num_anchors, 4) 。
# 生成的 bbox_mask 张量可以用于选择或过滤边界框的坐标。由于边界框通常由4个坐标（x1, y1, x2, y2）表示，因此 bbox_mask 的每个 num_anchors 行都会有4个元素，对应于边界框的每个坐标。
# 在这种情况下， bbox_mask 可以用于：
# 选择或过滤边界框的坐标，例如，在计算损失函数时只考虑被 fg_mask 标识为前景的边界框。
# 在某些操作中，如损失函数的计算，确保只有前景的边界框被考虑进去。
# 这种方法允许模型在训练过程中专注于那些更有可能包含目标物体的边界框，从而提高检测的准确性和效率。
bbox_mask = fg_mask.unsqueeze(-1).repeat([1, 1, 4])
# torch.masked_select(input, mask, out=None)
# torch.masked_select 是 PyTorch 中的一个函数，它用于根据布尔掩码（mask）从输入张量（input tensor）中选择元素。这个函数会返回一个一维张量，其中包含了所有掩码值为 True 对应的输入张量中的元素。
# input （Tensor）- 需要进行索引操作的输入张量。
# mask BoolTensor）- 布尔掩码，其形状必须与输入张量 input 通过广播机制兼容。掩码中的 True 表示选择 input 中对应的元素， False 表示不选择。
# out （Tensor, optional）- 指定输出的张量。如果提供，输出结果将被写入此张量。
# 返回值是一个一维张量，包含了所有 mask 中对应 True 的 input 元素。
pred_bboxes_pos = torch.masked_select(pred_bboxes,
bbox_mask).reshape([-1, 4])
target_bboxes_pos = torch.masked_select(
target_bboxes, bbox_mask).reshape([-1, 4])
# target_scores.sum(-1) ：沿着最后一个维度（ num_classes ）对 target_scores 中的元素进行求和。
bbox_weight = torch.masked_select(
target_scores.sum(-1), fg_mask).unsqueeze(-1)
# class IOUloss: -> 计算IoU损失。 -> def __call__(self, box1, box2): -> return loss
loss_iou = self.iou_loss(pred_bboxes_pos,
target_bboxes_pos) * bbox_weight
if target_scores_sum == 0:
loss_iou = loss_iou.sum()
else:
# 这种规范化操作的目的可能是为了计算损失函数的平均值，使得在不同批次大小或不同数量的正样本之间保持一致性。
# 在目标检测中，由于正样本（即包含目标物体的预测框）的数量可能在不同批次中变化，因此通过除以 target_scores_sum 可以确保损失函数的尺度不会因为正样本数量的不同而产生偏差。
# 最终， loss_iou 被规范化为一个标量值，这个值可以用于反向传播（backpropagation）来更新模型的权重。在训练过程中，这个平均损失值提供了一个衡量模型在当前批次上性能的指标。
loss_iou = loss_iou.sum() / target_scores_sum
# dfl loss 分布式焦点损失（Distributed Focal Loss，简称DFL）。
if self.use_dfl:
# 这段代码用于创建一个三维张量 dist_mask ，它将二维的 fg_mask 扩展到一个新的维度，并重复最后一个维度的元素。
# fg_mask.unsqueeze(-1) ： fg_mask 是一个二维张量，形状是 (batch_size, num_anchors) ，其中 num_anchors 是每个特征图单元格上的锚点（anchor）数量。
# unsqueeze(-1) 操作在张量的最后一个维度上增加一个维度，使其形状变为 (batch_size, num_anchors, 1) 。
# repeat([1, 1, (self.reg_max + 1) * 4]) ： repeat 操作将张量中的每个元素沿着指定的维度重复指定的次数。
# 在这里， [1, 1, (self.reg_max + 1) * 4] 表示在第一个和第二个维度（即 batch_siz 和 num_anchors ）上不重复，在最后一个维度（即新增加的维度）上重复 (self.reg_max + 1) * 4 次。
# 这样， dist_mask 的形状就变成了 (batch_size, num_anchors, (self.reg_max + 1) * 4) 。
# 生成的 dist_mask 张量可以用于选择或过滤边界框的坐标。由于边界框通常由 4 个坐标（x1, y1, x2, y2）表示，因此 dist_mask 的每个 num_anchors 行都会有 (self.reg_max + 1) * 4 个元素，对应于边界框的每个坐标。
# 这种方法允许模型在训练过程中专注于那些更有可能包含目标物体的边界框，从而提高检测的准确性和效率。
dist_mask = fg_mask.unsqueeze(-1).repeat(
[1, 1, (self.reg_max + 1) * 4])
# .reshape([-1, 4, self.reg_max + 1]) ： reshape 操作将一维张量重新排列成一个新的形状。 -1 表示自动计算该维度的大小，以确保总的元素数量保持不变。
# 这里，我们希望将张量重新排列成形状 (batch_size * num_anchors, 4, self.reg_max + 1) 。每个边界框有 4 个坐标（x1, y1, x2, y2），每个坐标有 self.reg_max + 1 个可能的值。
pred_dist_pos = torch.masked_select(
pred_dist, dist_mask).reshape([-1, 4, self.reg_max + 1])
# def bbox2dist(anchor_points, bbox, reg_max): -> 将 bbox(xyxy) 转换为 dist(ltrb)。 -> dist = torch.cat([lt, rb], -1).clip(0, reg_max - 0.01) -> return dist
target_ltrb = bbox2dist(anchor_points, target_bboxes, self.reg_max)
target_ltrb_pos = torch.masked_select(
target_ltrb, bbox_mask).reshape([-1, 4])
# def _df_loss(self, pred_dist, target): -> return (loss_left + loss_right).mean(-1, keepdim=True)
# loss_dfl 边界框回归的损失。
loss_dfl = self._df_loss(pred_dist_pos,
target_ltrb_pos) * bbox_weight
if target_scores_sum == 0:
loss_dfl = loss_dfl.sum()
else:
loss_dfl = loss_dfl.sum() / target_scores_sum
else:
loss_dfl = pred_dist.sum() * 0.
else:
loss_iou = pred_dist.sum() * 0.
loss_dfl = pred_dist.sum() * 0.
return loss_iou, loss_dfl
def _df_loss(self, pred_dist, target):
# target 是一个张量，它包含了目标物体边界框的左边缘的索引值。这些索引值被转换为长整型（ torch.long ），然后通过加 1 来得到右边缘的索引值。
# 将 target 张量的数据类型转换为长整型（ torch.long ）。这通常在处理索引时是必要的，因为索引需要是整数类型。
target_left = target.to(torch.long)
# target_left 的每个元素上加 1 来计算边界框右边缘的索引值。由于边界框的左边缘和右边缘在索引上是连续的，加 1 可以得到右边缘的索引。
# 这段代码的结果是两个张量， target_left 和 target_right ，它们分别包含了边界框左边缘和右边缘的索引值。这些索引值通常用于从预测的距离分布中选择相应的值，以便在损失函数中计算损失。
# 例如，如果 target 是一个形状为 (batch_size, num_anchors) 的张量，那么 target_left 和 target_right 也将是相同形状的张量，其中每个元素是对应边界框左边缘和右边缘的索引值。
target_right = target_left + 1
# 这个操作计算了每个边界框左边缘的权重。由于 target_right 是 target_left + 1 ，这个操作实际上是计算 1 - target ，其中 target 是左边缘的索引。这给出了每个边界框左边缘的权重。
weight_left = target_right.to(torch.float) - target
# 这个操作计算了每个边界框右边缘的权重。由于权重的总和应该是 1，所以右边缘的权重是 1 - weight_left 。
# 这段代码的结果是两个张量， weight_left 和 weight_right ，它们分别包含了边界框左边缘和右边缘的权重。这些权重通常用于在损失函数中平衡左边缘和右边缘的贡献，特别是在处理不平衡数据集或不同尺寸的边界框时。
# 例如，如果 target 是一个形状为 (batch_size, num_anchors) 的张量，那么 weight_left 和 weight_right 也将是相同形状的张量，其中每个元素是对应边界框左边缘和右边缘的权重。
# 这种方法有助于确保损失函数对边界框的每个部分都给予适当的关注，从而提高模型在回归任务中的性能。
weight_right = 1 - weight_left
# torch.nn.functional.cross_entropy(input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')
# torch.nn.functional.cross_entropy 是 PyTorch 中的一个函数，用于计算交叉熵损失（Cross-Entropy Loss），常用于多分类问题。这个函数结合了 softmax 操作和负对数似然损失（Negative Log Likelihood Loss, NLL Loss）的计算。
# 参数说明：
# input （Tensor）：模型输出的 logits，即未经 softmax 转换的原始预测值。形状通常是 (batch_size, num_classes) ，其中 batch_size 是批次大小， num_classes 是类别的数量。
# target （Tensor）：真实标签的索引值，形状为 (batch_size,) 。每个元素的值应该是 0 到 num_classes-1 之间的整数，表示对应样本的类别。
# weight （Tensor, optional）：每个类别的权重。如果指定，形状应该是 (num_classes,) ，用于对不同类别的损失进行加权。
# size_average （bool, optional）：在 PyTorch 0.4 之前的版本中用于指定是否对损失进行平均。在新版本中，这个参数已经被 reduction 参数替代。
# ignore_index （int, optional）：指定一个类别的索引，该类别的损失将被忽略（不会对梯度产生贡献）。通常用于处理不关心的类别。
# reduce （bool, optional）：在 PyTorch 0.4 之前的版本中用于指定是否对损失进行求和或平均。在新版本中，这个参数已经被 reduction 参数替代。
# reduction （string, optional）：指定损失的缩减方式。可以是 'none' （不缩减）， 'mean' （平均值），或者 'sum' （总和）。默认是 'mean' 。
# 函数的工作流程如下：
# 1. 对 input 张量进行 softmax 转换，得到每个样本属于每个类别的概率分布。
# 2. 使用 target 张量中的索引值，从 softmax 转换后的概率分布中选择对应类别的概率。
# 3. 计算这些概率的负对数，得到每个样本的损失值。
# 4. 根据 reduction 参数的设置，对所有样本的损失值进行平均或求和。
# 最终， cross_entropy 函数返回一个标量张量，表示整个批次的平均损失值（如果 reduction='mean' ），或者所有样本的损失总和（如果 reduction='sum' ），或者每个样本的损失值（如果 reduction='none' ）。这个值可以用于反向传播，以更新模型的权重。
loss_left = F.cross_entropy(
pred_dist.view(-1, self.reg_max + 1), target_left.view(-1), reduction='none').view(
target_left.shape) * weight_left
loss_right = F.cross_entropy(
pred_dist.view(-1, self.reg_max + 1), target_right.view(-1), reduction='none').view(
target_left.shape) * weight_right
# .mean(-1, keepdim=True) ： mean 函数计算张量在指定维度上的均值。 -1 表示最后一个维度， keepdim=True 表示在计算均值后保持原始张量的维度信息。
# 原始的 loss_left 和 loss_right 张量的形状是 (batch_size, num_anchors) ，那么相加后的形状仍然是 (batch_size, num_anchors) 。
# 在最后一个维度上计算均值后，结果的形状将是 (batch_size, 1) ，因为 keepdim=True 保持了原始的批次维度。
# 这个操作通常用于将不同部分的损失值（如边界框回归的左边缘和右边缘损失）合并成一个标量值，以便进一步处理，例如在损失函数中进行加权或直接用于反向传播。
return (loss_left + loss_right).mean(-1, keepdim=True)