loss_tal_triple.py
utils\loss_tal_triple.py
目录
loss_tal_triple.py
1.所需的库和模块
2.def smooth_BCE(eps=0.1):
3.class VarifocalLoss(nn.Module):
4.class FocalLoss(nn.Module):
5.class BboxLoss(nn.Module):
6.class ComputeLoss:
1.所需的库和模块
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.general import xywh2xyxy
from utils.metrics import bbox_iou
from utils.tal.anchor_generator import dist2bbox, make_anchors, bbox2dist
from utils.tal.assigner import TaskAlignedAssigner
from utils.torch_utils import de_parallel2.def smooth_BCE(eps=0.1):
# 这段代码定义了一个名为 smooth_BCE 的函数，它用于计算平滑的二元交叉熵损失（Binary Cross-Entropy, BCE）的边界值。在机器学习中，二元交叉熵损失函数常用于二分类问题，其中模型的输出是一个概率值，表示属于正类的概率。
# 定义了一个名为 smooth_BCE 的函数，它接受一个参数。
# 1.eps ：其默认值为 0.1 。 eps 是一个用于平滑的很小的正数，目的是避免在计算损失时出现数值不稳定的情况，比如除以零。
def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
# return positive, negative label smoothing BCE targets 返回正、负标签平滑 BCE 目标。
# 函数返回两个值，分别是 1.0 - 0.5 * eps 和 0.5 * eps 。这两个值分别对应于平滑后的二元交叉熵损失函数中的两个边界值。
# 在原始的二元交叉熵损失函数中，当预测值接近1时，损失接近0；当预测值接近0时，损失接近1。通过引入 eps ，我们可以将这两个边界值稍微移动，从而避免在极端情况下的数值问题。
return 1.0 - 0.5 * eps, 0.5 * eps
# smooth_BCE 函数是一个辅助函数，用于生成平滑二元交叉熵损失函数的两个边界值。通过设置 eps 参数，我们可以控制边界值的平滑程度，以提高数值稳定性。这个函数通常用于机器学习中，特别是在处理分类问题时，以避免在计算损失函数时出现极端值导致的数值问题。3.class VarifocalLoss(nn.Module):
# 这段代码定义了一个名为 VarifocalLoss 的类，它是一个继承自 nn.Module 的 PyTorch 模块，用于实现变分焦点损失（Varifocal Loss）。
# 定义了一个名为 VarifocalLoss 的类，它继承自 PyTorch 的 nn.Module 类，这是所有神经网络模块的基类。
class VarifocalLoss(nn.Module):
# Varifocal loss by Zhang et al. https://arxiv.org/abs/2008.13367 Zhang 等人的变焦损失 https://arxiv.org/abs/2008.13367 。
# 这是一个特殊的方法，当创建 VarifocalLoss 类的实例时会被自动调用。它初始化了类的实例。
def __init__(self):
# 调用了父类 nn.Module 的 __init__ 方法，这是初始化模块时的标准做法。
super().__init__()
# forward 方法定义了模块的前向传播逻辑。它接受四个参数。
# 1.pred_score ：预测分数。
# 2.gt_score ：真实分数。
# 3.label ：标签。
# 4.alpha 和 5.gamma （两个可选参数） ：它们分别默认为 0.75 和 2.0 。这些参数用于调整损失函数的行为。
def forward(self, pred_score, gt_score, label, alpha=0.75, gamma=2.0):
# 计算了每个样本的权重。权重是基于 预测分数 、 真实分数 和 标签 计算的。 alpha 和 gamma 是超参数，用于调整权重。 pred_score.sigmoid() 将预测分数通过 sigmoid 函数转换为概率值， pow(gamma) 则将这个概率值提升到 gamma 次幂。 (1 - label) 和 label 分别用于处理正样本和负样本。
weight = alpha * pred_score.sigmoid().pow(gamma) * (1 - label) + gt_score * label
# 使用 PyTorch 的自动混合精度（AMP）上下文管理器 autocast ，禁用自动混合精度。这是因为在计算损失时，需要使用全精度（float32）来避免数值不稳定。
with torch.cuda.amp.autocast(enabled=False):
# torch.nn.functional.binary_cross_entropy_with_logits(input, target, weight=None, pos_weight=None, reduction='mean')
# F.binary_cross_entropy_with_logits 是 PyTorch 中的一个函数，它计算二元交叉熵损失（Binary Cross Entropy Loss），这个损失函数适用于二分类问题。该函数结合了 Sigmoid 激活函数和二元交叉熵损失计算，使得它在数值上更加稳定，并且减少了计算步骤。
# 参数 ：
# input ：模型输出的 logits（即未经 Sigmoid 激活的原始输出），形状为 (N, *) ，其中 N 是批次大小， * 表示任意数量的附加维度。
# target ：真实标签，形状与 input 相同，值在 [0, 1] 范围内。
# weight ：每个样本的权重，可以用来处理不平衡数据集，形状为 (N, *) 。
# pos_weight ：正样本的权重，用于处理不平衡数据集中正样本较少的情况，形状为 (1,) 。
# reduction ：指定如何应用损失的缩减。可选的值为 'none' 、 'mean' 、 'sum' 。默认为 'mean' ，表示计算所有损失的平均值。
# 返回值 ：
# 返回一个标量或张量，取决于 reduction 参数的设置。
# 特点 ：
# 内部 Sigmoid ： F.binary_cross_entropy_with_logits 在计算损失之前，内部自动应用 Sigmoid 函数将 logits 转换为概率值，因此不需要在外部手动应用 Sigmoid。
# 数值稳定性 ：由于结合了 Sigmoid 和损失计算，该函数利用了 log-sum-exp 技巧来提高数值稳定性，这在处理极端值时尤为重要。
# 这个函数是处理二分类问题时的推荐选择，因为它减少了手动应用 Sigmoid 激活的步骤，并且提供了更好的数值稳定性。
# 计算了变分焦点损失。首先，它使用 F.binary_cross_entropy_with_logits 函数计算二元交叉熵损失，其中 reduction="none" 参数表示不对损失进行求和或平均，而是保留每个样本的损失值。然后，它将每个样本的损失与之前计算的权重相乘，最后对所有样本的加权损失求和。
loss = (F.binary_cross_entropy_with_logits(pred_score.float(), gt_score.float(),
reduction="none") * weight).sum()
# 返回计算得到的总损失。
return loss
# VarifocalLoss 类实现了变分焦点损失函数，这是一种用于目标检测任务的损失函数，特别适用于类别不平衡的情况。它通过调整权重来聚焦于难以分类的样本，从而提高模型的性能。这个类可以作为 PyTorch 模型训练过程中的一个损失函数模块。4.class FocalLoss(nn.Module):
# 这段代码定义了一个名为 FocalLoss 的类，它是一个包装类，用于将焦点损失（Focal Loss）应用于现有的损失函数。
# 定义了一个名为 FocalLoss 的类，它继承自 PyTorch 的 nn.Module 类。
class FocalLoss(nn.Module):
# Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 将焦点损失包裹在现有的 loss_fcn() 周围，即标准 = FocalLoss(nn.BCEWithLogitsLoss()，gamma=1.5) 。
# 这是 FocalLoss 类的构造函数，它接受三个参数。
# 1.loss_fcn ：一个损失函数。
# 2.gamma ：一个控制焦点损失形状的参数，默认值为1.5。
# 3.alpha ：一个平衡正负样本的参数，默认值为0.25。
def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
# 调用父类 nn.Module 的构造函数，是初始化模块时的标准做法。
super().__init__()
# 将传入的损失函数赋值给类的成员变量 loss_fcn 。
self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
# 将 gamma 参数赋值给类的成员变量 gamma 。
self.gamma = gamma
# 将 alpha 参数赋值给类的成员变量 alpha 。
self.alpha = alpha
# 将传入损失函数的 reduction 属性赋值给类的成员变量 reduction 。
self.reduction = loss_fcn.reduction
# 将传入损失函数的 reduction 属性设置为 "none" ，这是为了确保焦点损失可以应用于每个元素。
self.loss_fcn.reduction = "none" # required to apply FL to each element
# forward 方法定义了模块的前向传播逻辑。它接受两个参数。
# 1.pred ：预测值。
# 2.true ：真实值。
def forward(self, pred, true):
# 使用成员变量 loss_fcn 计算基本的损失。
loss = self.loss_fcn(pred, true)
# p_t = torch.exp(-loss)
# loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability
# TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
# 将 预测值 通过 sigmoid 函数转换为概率值。
pred_prob = torch.sigmoid(pred) # prob from logits
# 计算 p_t ，即 调整后的预测概率 ，它考虑了真实标签。
p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
# 计算 alpha_factor ，它根据 真实标签 调整 alpha 参数。
alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
# 计算 modulating_factor ，它根据 gamma 参数调整焦点损失。
modulating_factor = (1.0 - p_t) ** self.gamma
# 将基本损失与 alpha_factor 和 modulating_factor 相乘，得到焦点损失。
loss *= alpha_factor * modulating_factor
# 检查 reduction 属性是否为 "mean" ，如果是，则返回损失的平均值。
if self.reduction == "mean":
return loss.mean()
# 检查 reduction 属性是否为 "sum" ，如果是，则返回损失的总和。
elif self.reduction == "sum":
return loss.sum()
# 处理 reduction 属性为 "none" 的情况，返回未经进一步处理的损失。
else: # 'none'
return loss
# FocalLoss 类是一个包装器，它将焦点损失应用于基本的二元交叉熵损失函数。焦点损失通过调整损失函数来关注那些难以分类的样本，这在类别不平衡的情况下特别有用。这个类可以作为 PyTorch 模型训练过程中的一个损失函数模块。5.class BboxLoss(nn.Module):
# 这段代码定义了一个名为 BboxLoss 的类，它是一个继承自 nn.Module 的 PyTorch 模块，用于计算目标检测任务中的边界框损失。
# 定义了一个名为 BboxLoss 的类，它继承自 PyTorch 的 nn.Module 类。
class BboxLoss(nn.Module):
# 这是 BboxLoss 类的构造函数，它接受两个参数。
# 1.reg_max ：用于定义预测距离的最大值。
# 2.use_dfl ：一个布尔值，指示是否使用动态焦点损失，默认为False。
def __init__(self, reg_max, use_dfl=False):
# 调用父类 nn.Module 的构造函数，是初始化模块时的标准做法。
super().__init__()
# 将 reg_max 参数赋值给类的成员变量 reg_max 。
self.reg_max = reg_max
# 将 use_dfl 参数赋值给类的成员变量 use_dfl 。
self.use_dfl = use_dfl
# 这段代码是 BboxLoss 类的 forward 方法，它负责计算边界框损失。
# 这是 BboxLoss 类的 forward 方法，它定义了如何计算损失。这个方法接受七个参数。
# 1.pred_dist ：预测的距离。这个参数通常是一个张量，包含了模型对于每个锚点（anchor point）预测的边界框距离信息。这些距离信息可以用来调整锚点的位置，以更好地匹配目标边界框。
# 2.pred_bboxes ：预测的边界框。这个参数是一个张量，包含了模型预测的边界框坐标。这些坐标通常是相对于锚点的，需要进一步调整以匹配目标边界框。
# 3.anchor_points ：锚点。这是一个张量，包含了用于目标检测任务中的锚点坐标。锚点是预先定义的边界框，用于与目标边界框进行匹配。
# 4.target_bboxes ：目标边界框。这是一个张量，包含了真实目标的边界框坐标。这些坐标用于计算损失，以训练模型更准确地预测边界框。
# 5.target_scores ：目标分数。这是一个张量，包含了每个目标的存在概率分数。这些分数通常用于计算损失时的权重，以区分不同目标的重要性。
# 6.target_scores_sum ：目标分数总和。这是一个标量值，表示所有目标分数的总和。这个值用于在计算损失时进行归一化，以确保损失不会因为目标数量的不同而产生偏差。
# 7.fg_mask ：前景掩码。这是一个二进制张量，用于区分前景（目标）和背景。在目标检测任务中，前景掩码用于选择与目标相关的预测和目标信息，而忽略背景。
def forward(self, pred_dist, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask):
# iou loss
# 创建一个掩码 bbox_mask ，它将 fg_mask 在最后一个维度上扩展并重复四次，以便与边界框的四个坐标对应。
bbox_mask = fg_mask.unsqueeze(-1).repeat([1, 1, 4]) # (b, h*w, 4)
# torch.masked_select(input, mask, out=None)
# torch.masked_select 是 PyTorch 中的一个函数，它根据布尔掩码（mask）从输入张量中选择元素。
# 参数 ：
# input (Tensor) ：需要进行索引操作的输入张量。
# mask (BoolTensor) ：要进行索引的布尔掩码，其形状必须与输入张量 input 相同，或者可以广播到输入张量的形状。
# out (Tensor, optional) ：指定输出的张量。如果提供，输出结果将被存储在这个张量中，否则将创建一个新的张量来存储结果。
# 功能描述 ：
# torch.masked_select 函数返回一个一维张量，其中包含所有在 mask 中对应位置为 True 的 input 张量中的元素。如果 mask 中的元素为 False ，则对应的 input 中的元素不会被选中。
# 返回值 ：
# 一个一维张量，包含从输入张量中根据掩码选择的元素。
# 注意事项 ：
# 掩码 mask 必须是一个布尔张量，其数据类型为 torch.bool 。
# 掩码 mask 的形状必须与输入张量 input 相同，或者可以广播到输入张量的形状。
# 如果 out 参数被提供，其数据类型必须与 input 中被选中的元素的数据类型相同。
# torch.masked_select 在深度学习中应用广泛，特别是在处理需要根据条件选择数据的场景，如在计算损失函数时选择特定的预测结果，或者在数据预处理时过滤无效的数据等。
# 使用 bbox_mask 从 pred_bboxes 中选择前景（即被掩码标记为1的）边界框，并将其重塑为二维张量，其中每行代表一个边界框的四个坐标。
pred_bboxes_pos = torch.masked_select(pred_bboxes, bbox_mask).view(-1, 4)
# 同样，使用 bbox_mask 从 target_bboxes 中选择前景边界框，并将其重塑为二维张量。
target_bboxes_pos = torch.masked_select(target_bboxes, bbox_mask).view(-1, 4)
# 计算每个前景边界框的权重，通过对 target_scores 在最后一个维度上求和并使用 fg_mask 选择前景样本，然后扩展维度以匹配损失的形状。
bbox_weight = torch.masked_select(target_scores.sum(-1), fg_mask).unsqueeze(-1)
# 计算预测边界框和目标边界框之间的交并比（IoU），这里使用的是 CIoU（Complete IoU）方法，它是一种改进的 IoU 计算方法。
# def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, MDPIoU=False, feat_h=640, feat_w=640, eps=1e-7):
# -> 用于计算两个边界框之间的交并比（IoU）以及其变体，如GIoU、DIoU、CIoU和MDPIoU。
# -> return iou - (rho2 / c2 + v * alpha) # CIoU return iou - rho2 / c2 # DIoU return iou - (c_area - union) / c_area # GIoU return iou - d1 / mpdiou_hw_pow - d2 / mpdiou_hw_pow # MPDIoU return iou # IoU
iou = bbox_iou(pred_bboxes_pos, target_bboxes_pos, xywh=False, CIoU=True)
# 计算 IoU 损失，即 1 减去 IoU 值，这个损失衡量预测边界框与目标边界框之间的差异。
loss_iou = 1.0 - iou
# 将 IoU 损失与之前计算的权重相乘，以对不同的边界框应用不同的权重。
loss_iou *= bbox_weight
# 对加权的 IoU 损失求和，然后除以 target_scores_sum 以得到平均 IoU 损失。
loss_iou = loss_iou.sum() / target_scores_sum
# dfl loss
# 检查是否使用动态焦点损失（Dynamic Focal Loss，简称 DFL）。
if self.use_dfl:
# 如果使用 DFL，创建一个掩码 dist_mask ，它将 fg_mask 扩展并重复多次，以匹配预测距离的数量。
dist_mask = fg_mask.unsqueeze(-1).repeat([1, 1, (self.reg_max + 1) * 4])
# 使用 dist_mask 从 pred_dist 中选择前景边界框对应的预测距离，并将其重塑为三维张量。
pred_dist_pos = torch.masked_select(pred_dist, dist_mask).view(-1, 4, self.reg_max + 1)
# 将目标边界框转换为相对于锚点的距离形式。
# def bbox2dist(anchor_points, bbox, reg_max):
# -> 将边界框从 xyxy （左上角和右下角的坐标）格式转换为 ltrb （左、上、右、下的偏移量）格式。使用 torch.cat 将这两个偏移量沿着最后一个维度（即坐标维度）拼接起来，形成 ltrb 格式的距离。然后使用 clamp 函数将距离限制在 [0, reg_max - 0.01] 的范围内，以避免超出回归任务的最大值。
# -> return torch.cat((anchor_points - x1y1, x2y2 - anchor_points), -1).clamp(0, reg_max - 0.01) # dist (lt, rb)
target_ltrb = bbox2dist(anchor_points, target_bboxes, self.reg_max)
# 使用 bbox_mask 从 target_ltrb 中选择前景边界框对应的距离。
target_ltrb_pos = torch.masked_select(target_ltrb, bbox_mask).view(-1, 4)
# 计算动态焦点损失，并将其与之前计算的权重相乘。
loss_dfl = self._df_loss(pred_dist_pos, target_ltrb_pos) * bbox_weight
# 对加权的 DFL 求和，然后除以 target_scores_sum 以得到平均 DFL。
loss_dfl = loss_dfl.sum() / target_scores_sum
# 如果不使用 DFL，这行代码指定了 else 分支的开始。
else:
# 在不使用 DFL 的情况下，将 DFL 设置为0，并确保它与 pred_dist 在同一设备上。
loss_dfl = torch.tensor(0.0).to(pred_dist.device)
# 返回计算得到的 IoU 损失、DFL 和 IoU 值。
return loss_iou, loss_dfl, iou
# 这个方法结合了 IoU 损失和可选的 DFL，以评估预测边界框与目标边界框之间的差异，并根据预测的不确定性调整损失。
# 这段代码定义了一个名为 _df_loss 的私有方法，它是 BboxLoss 类的一个辅助方法，用于计算动态焦点损失（Dynamic Focal Loss，简称 DFL）。
# 这是 _df_loss 方法的定义，它接受两个参数。
# 1.pred_dist ：预测的距离。
# 2.target ：目标距离。
def _df_loss(self, pred_dist, target):
# 将目标距离向下取整到最近的整数，得到左边界。
target_left = target.to(torch.long)
# 计算右边界，即左边界加1。
target_right = target_left + 1
# 计算左边界的权重，即右边界值减去目标值。
weight_left = target_right.to(torch.float) - target
# 计算右边界的权重，即1减去左边界的权重。
weight_right = 1 - weight_left
# 计算左边界的交叉熵损失。首先，将预测距离重塑为一维，并计算交叉熵损失，其中 reduction="none" 表示不进行损失的求和或平均，保留每个样本的损失值。然后，将损失重塑为与 target_left 相同的形状，并乘以左边界的权重。
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
# 将左边界和右边界的损失相加，然后计算最后一个维度（即每个样本的损失）的平均值，并保持维度（ keepdim=True ）以便与原始损失张量的形状兼容。
return (loss_left + loss_right).mean(-1, keepdim=True)
# _df_loss 方法计算了每个预测距离的动态焦点损失，考虑了目标距离的左边界和右边界，并为每个边界计算了加权的交叉熵损失。这种方法有助于模型更精确地学习边界框的位置，特别是在目标边界框与锚点距离较近时。最终返回的是每个样本的平均 DFL，可以用于进一步的损失计算。
# BboxLoss 类实现了目标检测任务中的边界框损失计算，包括IoU损失和可选的动态焦点损失。这个类可以作为 PyTorch 模型训练过程中的一个损失函数模块，用于优化目标检测模型的性能。6.class ComputeLoss:
# 这段代码定义了一个名为 ComputeLoss 的类，它用于计算目标检测模型的损失。
# 定义了一个名为 ComputeLoss 的类，用于计算目标检测模型的损失。
class ComputeLoss:
# Compute losses
# 这段代码是 ComputeLoss 类的构造函数 __init__ ，它初始化类实例并设置计算损失所需的参数和组件。
# 这是 ComputeLoss 类的构造函数，它接受两个参数。
# 1.model ：目标检测模型。
# 2.use_dfl ：一个布尔值，指示是否使用动态焦点损失，默认为True。
def __init__(self, model, use_dfl=True):
# 获取模型参数所在的设备（CPU或GPU），以便后续将损失函数和张量放置在相同的设备上。
device = next(model.parameters()).device # get model device
# 获取模型的超参数，通常这些参数用于控制模型的训练和行为。
h = model.hyp # hyperparameters
# Define criteria
# 定义了一个二元交叉熵损失函数 BCEcls ，它用于分类任务。 pos_weight 参数用于平衡正负样本的数量， reduction='none' 表示不对损失进行求和或平均，以便可以对每个样本的损失进行操作。
BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device), reduction='none')
# Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
# 调用 smooth_BCE 函数来生成平滑的二元交叉熵损失的目标值，用于标签平滑技术， eps 参数控制平滑的程度。
# def smooth_BCE(eps=0.1): -> 用于计算平滑的二元交叉熵损失（Binary Cross-Entropy, BCE）的边界值。 -> return 1.0 - 0.5 * eps, 0.5 * eps
self.cp, self.cn = smooth_BCE(eps=h.get("label_smoothing", 0.0)) # positive, negative BCE targets
# Focal loss
# 从超参数中获取焦点损失的 gamma 值，它用于调整焦点损失的形状。
g = h["fl_gamma"] # focal loss gamma
# 检查 gamma 值是否大于0，如果是，则启用焦点损失。
if g > 0:
# 如果启用焦点损失，将 BCEcls 包装在焦点损失中，以关注那些难以分类的样本。
BCEcls = FocalLoss(BCEcls, g)
# 获取模型的最后一个检测模块， de_parallel 函数可能用于移除模型的并行封装。
# def de_parallel(model):
# -> 将一个可能处于并行状态（例如使用 PyTorch 的 DataParallel 或 DistributedDataParallel 包装过的模型）转换回单个 GPU 或 CPU 上的模型。如果 model 是并行模型， 将返回原始的、未并行化的模型 model.module 。如果 model 不是并行模型，直接返回 model 。
# -> return model.module if is_parallel(model) else model
m = de_parallel(model).model[-1] # Detect() module
# 设置不同层的损失权重平衡， m.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
# 将定义的二元交叉熵损失函数保存为类的成员变量。
self.BCEcls = BCEcls
# 将模型的超参数保存为类的成员变量。
self.hyp = h
# 保存模型的步长（stride），它用于将特征图的坐标映射回原始图像的坐标。
self.stride = m.stride # model strides
# 保存模型的类别数。
self.nc = m.nc # number of classes
# 保存模型的层数。
self.nl = m.nl # number of layers
# 保存模型输出的每个锚点的数量。
self.no = m.no
# 保存模型的最大回归值。
self.reg_max = m.reg_max
# 保存模型参数所在的设备。
self.device = device
# 创建一个任务对齐分配器实例，用于为每个预测分配真实目标。
self.assigner = TaskAlignedAssigner(topk=int(os.getenv('YOLOM', 10)),
num_classes=self.nc,
alpha=float(os.getenv('YOLOA', 0.5)),
beta=float(os.getenv('YOLOB', 6.0)))
# 创建第二个任务对齐分配器实例。
self.assigner2 = TaskAlignedAssigner(topk=int(os.getenv('YOLOM', 10)),
num_classes=self.nc,
alpha=float(os.getenv('YOLOA', 0.5)),
beta=float(os.getenv('YOLOB', 6.0)))
# 创建第三个任务对齐分配器实例。
self.assigner3 = TaskAlignedAssigner(topk=int(os.getenv('YOLOM', 10)),
num_classes=self.nc,
alpha=float(os.getenv('YOLOA', 0.5)),
beta=float(os.getenv('YOLOB', 6.0)))
# 创建一个 BboxLoss 实例，用于计算边界框损失，并将其移动到正确的设备上。
self.bbox_loss = BboxLoss(m.reg_max - 1, use_dfl=use_dfl).to(device)
# 创建第二个 BboxLoss 实例。
self.bbox_loss2 = BboxLoss(m.reg_max - 1, use_dfl=use_dfl).to(device)
# 创建第三个 BboxLoss 实例。
self.bbox_loss3 = BboxLoss(m.reg_max - 1, use_dfl=use_dfl).to(device)
# torch.arange(start=0, end=0, step=1, out=None, dtype=None, device=None, requires_grad=False) -> Tensor
# torch.arange 是 PyTorch 中的一个函数，它返回一个由连续整数组成的一维张量，类似于 Python 的内置函数 range 。这个函数通常用于创建序列或索引。
# 参数说明 ：
# start ：序列的起始值，默认为 0。
# end ：序列的结束值，但不包含此值。
# step ：序列中每个元素的步长，默认为 1。
# out ：一个可选的 Tensor ，用于存储输出结果。
# dtype ：输出张量的所需数据类型，默认为 torch.float32 。
# device ：输出张量的所需设备（CPU 或 GPU），默认为 CPU。
# requires_grad ：是否需要计算梯度，默认为 False。
# 返回值 ：
# 返回一个一维 Tensor ，包含了从 start 到 end （不包含）的整数序列，步长为 step 。
# 创建一个投影张量，用于DFL损失，并将其移动到正确的设备上。
self.proj = torch.arange(m.reg_max).float().to(device) # / 120.0
# 这行代码保存是否使用DFL的标志。
self.use_dfl = use_dfl
# ComputeLoss 类的构造函数初始化了目标检测模型损失计算所需的所有组件，包括分类损失、边界框损失和任务对齐分配器。它还处理了模型超参数的获取和设备的选择。这个类的实例将用于在训练过程中计算模型的损失。
# 这段代码定义了 ComputeLoss 类中的 preprocess 方法，该方法用于预处理目标检测任务中的目标信息。
# 这是 preprocess 方法的定义，它接受三个参数。
# 1.targets ：包含所有目标的信息。
# 2.batch_size ：当前批次的大小。
# 3.scale_tensor ：用于缩放目标边界框的张量。
def preprocess(self, targets, batch_size, scale_tensor):
# 检查 targets 是否为空（即没有目标）。
if targets.shape[0] == 0:
# 如果 targets 为空，则创建一个形状为 (batch_size, 0, 5) 的全零张量，其中 5 表示每个目标的信息（通常是 类别 和 四个边界框坐标 ）。
out = torch.zeros(batch_size, 0, 5, device=self.device)
# 如果 targets 不为空，执行以下操作。
else:
# 提取 targets 的第一列，它包含图像索引，用于将目标分配给相应的图像。
i = targets[:, 0] # image index
# 计算每个图像索引出现的次数， counts.max() 得到最大目标数，用于初始化输出张量的大小。
_, counts = i.unique(return_counts=True)
# 创建一个形状为 (batch_size, max_counts, 5) 的全零张量，用于存储预处理后的目标信息。
out = torch.zeros(batch_size, counts.max(), 5, device=self.device)
# 遍历每个批次中的图像。
for j in range(batch_size):
# 对于每个图像，找到所有匹配的目标。
matches = i == j
# 计算匹配目标的数量。
n = matches.sum()
# 如果存在匹配的目标，执行以下操作。
if n:
# 将匹配的目标信息复制到输出张量 out 中。
out[j, :n] = targets[matches, 1:]
# 将输出张量 out 中的边界框坐标从 (x, y, w, h) 格式转换为 (x1, y1, x2, y2) 格式，并应用缩放。
out[..., 1:5] = xywh2xyxy(out[..., 1:5].mul_(scale_tensor))
# 返回预处理后的目标信息。
return out
# preprocess 方法的主要作用是将目标信息整理成模型训练时所需的格式，并根据批次大小和图像尺寸对目标边界框进行缩放。这个方法确保了目标信息与模型输出的对应关系，并为后续的损失计算做好了准备。
# 这段代码定义了 ComputeLoss 类中的 bbox_decode 方法，该方法用于将预测的距离转换为边界框坐标。
# 这是 bbox_decode 方法的定义，它接受两个参数。
# anchor_points ：锚点坐标）和 pred_dist （预测的距离）。
def bbox_decode(self, anchor_points, pred_dist):
# 检查是否使用动态焦点损失（Dynamic Focal Loss，DFL）。
if self.use_dfl:
# 如果使用DFL，首先获取 pred_dist 的形状，分别代表 批次大小（ b ）、锚点数量（ a ）和通道数（ c ）。
b, a, c = pred_dist.shape # batch, anchors, channels
# 将 pred_dist 重塑为 (batch, anchors, 4, channels // 4) 的形状，并在最后一个维度上应用softmax函数，以获得每个锚点的预测分布。 将经过softmax处理的 pred_dist 与投影张量 self.proj 相乘，以获得最终的预测距离。
# self.proj ：一个从0到 reg_max 的整数张量，用于将预测的距离投影到实际的距离值。
pred_dist = pred_dist.view(b, a, 4, c // 4).softmax(3).matmul(self.proj.type(pred_dist.dtype))
# pred_dist = pred_dist.view(b, a, c // 4, 4).transpose(2,3).softmax(3).matmul(self.proj.type(pred_dist.dtype))
# pred_dist = (pred_dist.view(b, a, c // 4, 4).softmax(2) * self.proj.type(pred_dist.dtype).view(1, 1, -1, 1)).sum(2)
# 使用 dist2bbox 函数将预测的距离转换为边界框坐标。 xywh=False 参数表示输出的边界框格式为 (x1, y1, x2, y2) 而不是 (x, y, w, h) 。
return dist2bbox(pred_dist, anchor_points, xywh=False)
# bbox_decode 方法的作用是将模型预测的距离转换为边界框坐标。如果启用了DFL，它会先对预测的距离应用softmax函数，然后与投影张量相乘以获得最终的距离，最后将这些距离转换为边界框坐标。这个方法是目标检测模型中将预测转换为可用于计算损失的边界框的关键步骤。
# 这段代码定义了 ComputeLoss 类的 __call__ 方法，它是类的主方法，用于计算给定模型预测 p 和目标 targets 的损失。
# 这是 __call__ 方法的定义，它接受四个参数。
# 1.p ：模型预测。
# 2.targets ：目标。
# 3.img ：可选的图像。
# 4.epoch ：当前训练周期。
def __call__(self, p, targets, img=None, epoch=0):
# 这段代码是 __call__ 方法的一部分，它负责准备和处理模型输出以及目标检测任务中的特征图。
# 初始化一个包含三个元素的零张量，分别用于存储边 界框损失 （ box ） 、 分类损失 （ cls ） 和 动态焦点损失 （ dfl ）。
loss = torch.zeros(3, device=self.device) # box, cls, dfl
# 根据模型输出 p 的类型，提取第一层特征图 feats 。如果 p 是元组，则 feats 是 p[1][0] ；否则， feats 是 p[0] 。
feats = p[1][0] if isinstance(p, tuple) else p[0]
# 提取第二层特征图 feats2 。
feats2 = p[1][1] if isinstance(p, tuple) else p[1]
# 提取第三层特征图 feats3 。
feats3 = p[1][2] if isinstance(p, tuple) else p[2]
# 将 feats 中的每个特征图的输出在第二维连接，并分割为预测的距离（ pred_distri ）和分数（ pred_scores ）。 self.reg_max * 4 是距离参数的数量， self.nc 是类别数量。
pred_distri, pred_scores = torch.cat([xi.view(feats[0].shape[0], self.no, -1) for xi in feats], 2).split(
(self.reg_max * 4, self.nc), 1) # ❌ ⚠️ 不应该在维度1上进行分割，正确的应该是 在 形状为 (batch_size, self.no, total_elements)张量的 total_elements 维度，即维度2上进行分割。
# 调整 预测分数 的维度，使其变为 (batch_size, num_classes, num_anchors) 的形状，并确保内存连续。
pred_scores = pred_scores.permute(0, 2, 1).contiguous()
# 调整 预测距离 的维度，并确保内存连续。
pred_distri = pred_distri.permute(0, 2, 1).contiguous()
# 对第二层特征图 feats2 执行相同的操作，得到 pred_distri2 和 pred_scores2 。
pred_distri2, pred_scores2 = torch.cat([xi.view(feats2[0].shape[0], self.no, -1) for xi in feats2], 2).split(
(self.reg_max * 4, self.nc), 1) # ❌ ⚠️ 不应该在维度1上进行分割，正确的应该是 在 形状为 (batch_size, self.no, total_elements)张量的 total_elements 维度，即维度2上进行分割。
pred_scores2 = pred_scores2.permute(0, 2, 1).contiguous()
pred_distri2 = pred_distri2.permute(0, 2, 1).contiguous()
# 对第三层特征图 feats3 执行相同的操作，得到 pred_distri3 和 pred_scores3 。
pred_distri3, pred_scores3 = torch.cat([xi.view(feats3[0].shape[0], self.no, -1) for xi in feats3], 2).split(
(self.reg_max * 4, self.nc), 1) # ❌ ⚠️ 不应该在维度1上进行分割，正确的应该是 在 形状为 (batch_size, self.no, total_elements)张量的 total_elements 维度，即维度2上进行分割。
pred_scores3 = pred_scores3.permute(0, 2, 1).contiguous()
pred_distri3 = pred_distri3.permute(0, 2, 1).contiguous()
# 获取预测分数的 数据类型 ，以便后续操作中使用相同的数据类型。
dtype = pred_scores.dtype
# 从预测分数的形状中获取 批次大小 和 网格大小 。
batch_size, grid_size = pred_scores.shape[:2]
# 计算 图像尺寸 ，将特征图的空间维度（高度和宽度）乘以步长 self.stride[0] 。
imgsz = torch.tensor(feats[0].shape[2:], device=self.device, dtype=dtype) * self.stride[0] # image size (h,w)
# 生成锚点和步长张量， make_anchors 函数根据特征图、步长和锚点的宽度因子（这里是0.5）生成锚点。
# def make_anchors(feats, strides, grid_cell_offset=0.5): -> 用于生成YOLO模型中使用的锚点（anchor points）。使用 torch.cat 函数将 anchor_points 和 stride_tensor 列表中的所有张量连接起来，并返回结果。 -> return torch.cat(anchor_points), torch.cat(stride_tensor)
anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5)
# 这段代码负责处理模型输出，将其转换为适合损失计算的形式，并准备必要的信息，如锚点和图像尺寸。这些步骤是目标检测中损失计算的基础，确保了后续损失计算的准确性和有效性。
# 这段代码继续 __call__ 方法的处理流程，专注于对目标（ targets ）进行预处理，并从中提取标签和边界框信息。
# targets
# 调用 preprocess 方法对目标进行预处理。 preprocess 方法将目标信息调整为模型输出的格式，并根据图像尺寸 imgsz 对边界框进行缩放。 scale_tensor 参数是一个张量，包含了图像的宽度和高度，用于将边界框坐标缩放到原始图像尺寸。
targets = self.preprocess(targets, batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
# targets 张量被分割成两个部分。 gt_labels 和 gt_bboxes 。 split 方法的第一个参数是一个元组 (1, 4) ，表示在第二个维度（索引为2）上分割张量，第一个分割部分的大小为1（对应类别标签），第二个分割部分的大小为4（对应边界框坐标）。 2 参数指定了分割的维度。
gt_labels, gt_bboxes = targets.split((1, 4), 2) # cls, xyxy
# 计算一个掩码 mask_gt ，用于标记非空的目标边界框。 gt_bboxes.sum(2, keepdim=True) 在第三个维度（边界框坐标维度）上计算每个目标的边界框坐标之和，并保持维度不变。 gt_bboxes.sum(2, keepdim=True).gt_(0) 检查这个和是否大于0，从而确定每个目标是否有有效的边界框坐标（即非空目标）。
mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0)
# 这段代码处理目标信息，将其转换为适合损失计算的格式，并提取出类别标签和边界框坐标。通过计算掩码 mask_gt ，它还帮助识别哪些目标是有效的，这对于后续的损失计算至关重要，因为只有有效的目标才会被用于计算损失。
# 这段代码处理预测的边界框（pboxes）并使用分配器（assigner）将预测与真实目标进行匹配。
# pboxes
# 使用 bbox_decode 方法将 预测的距离 pred_distri 和 锚点 anchor_points 转换为 预测的边界框 pred_bboxes 。输出的边界框格式为 xyxy ，即左上角和右下角的坐标。
pred_bboxes = self.bbox_decode(anchor_points, pred_distri) # xyxy, (b, h*w, 4)
# 对第二层特征图的预测进行相同的处理，得到 pred_bboxes2 。
pred_bboxes2 = self.bbox_decode(anchor_points, pred_distri2) # xyxy, (b, h*w, 4)
# 对第三层特征图的预测进行相同的处理，得到 pred_bboxes3 。
pred_bboxes3 = self.bbox_decode(anchor_points, pred_distri3) # xyxy, (b, h*w, 4)
# 使用分配器 assigner 将预测分数、预测边界框和锚点与真实目标进行匹配。输出包括 目标标签 target_labels 、 目标边界框 target_bboxes 、 目标分数 target_scores 和 前景掩码 fg_mask 。
target_labels, target_bboxes, target_scores, fg_mask = self.assigner(
pred_scores.detach().sigmoid(),
(pred_bboxes.detach() * stride_tensor).type(gt_bboxes.dtype),
anchor_points * stride_tensor,
gt_labels,
gt_bboxes,
mask_gt)
# 对第二层特征图的预测执行相同的匹配操作，得到相应的 target_labels2 、 target_bboxes2 、 target_scores2 和 fg_mask2 。
target_labels2, target_bboxes2, target_scores2, fg_mask2 = self.assigner2(
pred_scores2.detach().sigmoid(),
(pred_bboxes2.detach() * stride_tensor).type(gt_bboxes.dtype),
anchor_points * stride_tensor,
gt_labels,
gt_bboxes,
mask_gt)
# 对第三层特征图的预测执行相同的匹配操作，得到相应的 target_labels3 、 target_bboxes3 、 target_scores3 和 fg_mask3 。
target_labels3, target_bboxes3, target_scores3, fg_mask3 = self.assigner3(
pred_scores3.detach().sigmoid(),
(pred_bboxes3.detach() * stride_tensor).type(gt_bboxes.dtype),
anchor_points * stride_tensor,
gt_labels,
gt_bboxes,
mask_gt)
# 将目标边界框 target_bboxes 除以步长张量 stride_tensor ，以将其从特征图坐标转换回原始图像坐标。❌
# ⚠️ 从上面self.assigner()得到的 真实边界框 target_bboxes 是在原始图像坐标系下的坐标，将它除以 步长 stride_tensor 得到的是在特征图坐标系下的坐标。上下文中从self.bbox_decode()得到的 预测边界框 pred_bboxes 它也是特征图坐标系下的坐标。
# ⚠️ 这两个坐标将用于下文self.bbox_loss()中计算损失。
# ⚠️ 在计算损失时涉及真实框与预测框的交并比（IoU）的计算。在计算交并比（IoU）时，真实框和预测框的坐标确实需要在同一坐标系下。通常情况下，这些坐标应该是相对于原始图像的绝对坐标，而不是特征图坐标系下的坐标。
# ⚠️ 我认为应该将 target_bboxes /= stride_tensor 代码改为 pred_bboxes *= stride_tensor ，即将 预测边界框 pred_bboxes 从特征图坐标系转换为原始坐标系，与本就是原始坐标系下的 真实边界框 target_bboxes 参与self.bbox_loss()中损失的计算。
target_bboxes /= stride_tensor
# 计算目标分数 target_scores 的总和，并确保至少为1，以避免除以零。
target_scores_sum = max(target_scores.sum(), 1)
# 将第二层特征图的目标边界框 target_bboxes2 除以步长张量 stride_tensor 。❌
target_bboxes2 /= stride_tensor
# 计算第二层特征图的目标分数 target_scores2 的总和，并确保至少为1。
target_scores_sum2 = max(target_scores2.sum(), 1)
# 将第三层特征图的目标边界框 target_bboxes3 除以步长张量 stride_tensor 。❌
target_bboxes3 /= stride_tensor
# 计算第三层特征图的目标分数 target_scores3 的总和，并确保至少为1。
target_scores_sum3 = max(target_scores3.sum(), 1)
# 这段代码负责将模型的预测转换为边界框，并使用分配器将这些预测与真实目标匹配。通过这个过程，得到了每个预测应该关注的目标标签和边界框，以及用于损失计算的前景掩码。这些信息对于计算分类损失和边界框损失至关重要。
# 这段代码完成了 分类损失 （ cls loss ）和 边界框损失 （ bbox loss ）的计算，并应用了不同的权重来平衡不同部分的损失。
# cls loss
# loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum # VFL way
# 计算第一层特征图的 分类损失 ，使用二元交叉熵损失函数 BCEcls ，并将结果乘以0.25后求和，然后除以 target_scores_sum 来进行归一化。
loss[1] = 0.25 * self.BCEcls(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum # BCE
# 对第二层特征图做同样的 分类损失 计算，并加到 loss[1] 上。
loss[1] += 0.25 * self.BCEcls(pred_scores2, target_scores2.to(dtype)).sum() / target_scores_sum2 # BCE
# 对第三层特征图做 分类损失 计算，不乘以0.25，直接加到 loss[1] 上。
loss[1] += self.BCEcls(pred_scores3, target_scores3.to(dtype)).sum() / target_scores_sum3 # BCE
# bbox loss
# 检查前景掩码 fg_mask 中是否有任何真值（即是否有前景目标）。
if fg_mask.sum():
# 如果有前景目标，计算第一层特征图的 边界框损失 和 动态焦点损失 ，以及 IoU分数 。
loss[0], loss[2], iou = self.bbox_loss(pred_distri,
pred_bboxes,
anchor_points,
target_bboxes,
target_scores,
target_scores_sum,
fg_mask)
# 将第一层特征图的 边界框损失 乘以0.25。
loss[0] *= 0.25
# 将第一层特征图的 动态焦点损失 乘以0.25
loss[2] *= 0.25
# 检查第二层特征图的前景掩码 fg_mask2 。
if fg_mask2.sum():
# 计算第二层特征图的边界框损失和动态焦点损失。
loss0_, loss2_, iou2 = self.bbox_loss2(pred_distri2,
pred_bboxes2,
anchor_points,
target_bboxes2,
target_scores2,
target_scores_sum2,
fg_mask2)
# 将第二层特征图的边界框损失乘以0.25后加到 loss[0] 上。
loss[0] += 0.25 * loss0_
# 将第二层特征图的动态焦点损失乘以0.25后加到 loss[2] 上。
loss[2] += 0.25 * loss2_
# 检查第三层特征图的前景掩码 fg_mask3 。
if fg_mask3.sum():
# 计算第三层特征图的边界框损失和动态焦点损失。
loss0__, loss2__, iou3 = self.bbox_loss3(pred_distri3,
pred_bboxes3,
anchor_points,
target_bboxes3,
target_scores3,
target_scores_sum3,
fg_mask3)
# 将第三层特征图的边界框损失加到 loss[0] 上。
loss[0] += loss0__
# 将第三层特征图的动态焦点损失加到 loss[2] 上。
loss[2] += loss2__
# 对边界框损失进行缩放，乘以7.5作为最终的 边界框损失权重。
loss[0] *= 7.5 # box gain
# 对分类损失进行缩放，乘以0.5作为最终的 分类损失权重。
loss[1] *= 0.5 # cls gain
# 对动态焦点损失进行缩放，乘以1.5作为最终的 动态焦点损失权重。
loss[2] *= 1.5 # dfl gain
# 计算总损失，将每个部分的损失相加，乘以批次大小，返回总损失和损失的副本（用于不影响梯度回传的中间结果）。
return loss.sum() * batch_size, loss.detach() # loss(box, cls, dfl)
# 这段代码综合了分类损失和边界框损失的计算，对不同层级的特征图损失进行了加权和缩放，以平衡不同损失对总损失的贡献。最终返回的是缩放后的总损失，用于模型的训练优化。
# __call__ 方法是 ComputeLoss 类的核心，它整合了模型预测、目标分配、损失计算等多个步骤，最终返回总损失和分类、边界框、DFL损失的组合。这个方法确保了损失计算的完整性和准确性，是训练目标检测模型的关键部分。
# 这个类包含了计算目标检测模型损失所需的所有组件，包括分类损失、边界框损失和DFL损失。它还处理了目标的预处理和预测的解码。
