ultralytics\utils\loss.py
目录
loss.py
1.所需的库和模块
2.class VarifocalLoss(nn.Module):
3.class FocalLoss(nn.Module):
4.class BboxLoss(nn.Module):
5.class RotatedBboxLoss(BboxLoss):
6.class KeypointLoss(nn.Module):
7.class v8DetectionLoss:
8.class v8SegmentationLoss(v8DetectionLoss):
9.class v8PoseLoss(v8DetectionLoss):
10.class v8ClassificationLoss:
11.class v8OBBLoss(v8DetectionLoss):
12.class v10DetectLoss:
# Ultralytics YOLO , AGPL-3.0 license
import torch
import torch.nn as nn
import torch.nn.functional as F
from ultralytics.utils.metrics import OKS_SIGMA
from ultralytics.utils.ops import crop_mask, xywh2xyxy, xyxy2xywh
from ultralytics.utils.tal import RotatedTaskAlignedAssigner, TaskAlignedAssigner, dist2bbox, dist2rbox, make_anchors
from .metrics import bbox_iou, probiou
from .tal import bbox2dist
# 这段代码定义了一个名为 VarifocalLoss 的类,它继承自PyTorch的 nn.Module 。 VarifocalLoss 类是一个用于计算变焦损失(Varifocal Loss)的 PyTorch 模块。变焦损失是一种用于目标检测任务中的损失函数,旨在更好地平衡正负样本之间的权重。
# 定义了一个名为 VarifocalLoss 的类,继承自 nn.Module 。 nn.Module 是 PyTorch 中所有神经网络模块的基类。
class VarifocalLoss(nn.Module):
# 张等人的变焦损失 https://arxiv.org/abs/2008.13367。
"""
Varifocal loss by Zhang et al.
https://arxiv.org/abs/2008.13367.
"""
# 定义了类的初始化方法。
def __init__(self):
# 初始化 VarifocalLoss 类。
"""Initialize the VarifocalLoss class."""
# 调用父类 nn.Module 的初始化方法。这一步是必要的,因为它确保了 VarifocalLoss 类正确地继承了 nn.Module 的所有属性和方法。
super().__init__()
# 这段代码定义了 VarifocalLoss 类的 forward 方法,用于计算变焦损失(Varifocal Loss)。
# @staticmethod 表示 forward 是一个静态方法,不需要实例化对象即可调用。
@staticmethod
# 定义了 forward 方法,接受以下参数 :
# 1.pred_score :预测得分,通常是模型输出的 logits。
# 2.gt_score :真实得分,通常是目标标签。
# 3.label :标签张量,用于指示每个样本是正样本还是负样本。
# 4.alpha 和 5.gamma :是超参数,用于调整损失函数的权重和形状,默认值分别为 0.75 和 2.0。
def forward(pred_score, gt_score, label, alpha=0.75, gamma=2.0):
# 计算机变焦损失。
"""Computes varfocal loss."""
# 计算了损失函数的权重 weight 。
# pred_score.sigmoid() 将预测得分通过 sigmoid 函数转换为概率值。
# pred_score.sigmoid().pow(gamma) 将转换后的概率值提升到 gamma 次幂,用于调整负样本的权重。
# alpha * pred_score.sigmoid().pow(gamma) * (1 - label) 计算负样本的权重,其中 1 - label 表示负样本的掩码。
# gt_score * label 计算正样本的权重,其中 label 表示正样本的掩码。
# 最后,将正负样本的权重相加得到最终的权重 weight 。
weight = alpha * pred_score.sigmoid().pow(gamma) * (1 - label) + gt_score * label
# 这个上下文管理器用于禁用自动混合精度(AMP)。在某些情况下,可能需要确保损失计算在全精度下进行,以避免精度问题。
with torch.cuda.amp.autocast(enabled=False):
# 计算二元交叉熵损失。 F 是 torch.nn.functional 的别名, binary_cross_entropy_with_logits 是一个函数,用于计算预测得分和真实得分之间的二元交叉熵损失。 reduction="none" 表示不对损失进行平均或求和,返回每个元素的损失值。
# * weight 将计算得到的损失值乘以之前计算的权重 weight 。
loss = (
# 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(pred_score.float(), gt_score.float(), reduction="none") * weight)
# .mean(1) 对每个样本的损失值求平均, 1 表示沿着第1个维度(通常是特征维度)进行操作。
.mean(1)
# sum() 对所有样本的平均损失值求和,得到最终的损失值 loss 。
.sum()
)
# 返回计算得到的损失值 loss 。
return loss
# VarifocalLoss 的 forward 方法通过计算加权的二元交叉熵损失来优化模型的分类性能。负样本的权重通过预测得分的 sigmoid 函数和 gamma 次幂进行调整,而正样本的权重则直接使用真实得分。最终的损失值是加权的二元交叉熵损失的平均值的总和。这种损失函数有助于在不平衡的数据集中更好地平衡正负样本的贡献,从而提高模型的分类准确性。
# VarifocalLoss 类实现了一种用于目标检测分类任务的损失函数。它通过计算正负样本的权重,并结合二元交叉熵损失,来优化模型的分类性能。负样本的权重通过预测得分的 sigmoid 函数和 gamma 次幂进行调整,而正样本的权重则直接使用真实得分。最终的损失值是加权的二元交叉熵损失的平均值的总和。这种损失函数有助于在不平衡的数据集中更好地平衡正负样本的贡献。
# 这段代码定义了一个名为 FocalLoss 的类,继承自 PyTorch 的 nn.Module ,用于计算焦点损失(Focal Loss)。焦点损失是一种用于解决类别不平衡问题的损失函数,特别适用于目标检测任务。
# 定义了一个名为 FocalLoss 的类,继承自 nn.Module 。
class FocalLoss(nn.Module):
# 将焦点损失包裹在现有的 loss_fcn() 周围,即标准 = FocalLoss(nn.BCEWithLogitsLoss(),gamma=1.5)。
"""Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)."""
# 定义了类的初始化方法。
def __init__(self):
# 没有参数的 FocalLoss 类的初始化程序。
"""Initializer for FocalLoss class with no parameters."""
# 调用父类 nn.Module 的初始化方法,确保 FocalLoss 类正确地继承了 nn.Module 的所有属性和方法。
super().__init__()
# 这段代码是 FocalLoss 类的 forward 方法的实现,用于计算焦点损失(Focal Loss)。
# @staticmethod 表示 forward 是一个静态方法,不需要实例化对象即可调用。
@staticmethod
# 定义了 forward 方法,接受以下参数 :
# 1.pred :预测的 logits,通常是模型输出的原始得分。
# 2.label :标签张量,表示每个样本的真实类别。
# 3.gamma 和 4.alpha :是超参数,分别默认为 1.5 和 0.25,用于调整损失函数的权重和形状。
def forward(pred, label, gamma=1.5, alpha=0.25):
# 计算并更新对象检测/分类任务的混淆矩阵。
"""Calculates and updates confusion matrix for object detection/classification tasks."""
# 使用 F.binary_cross_entropy_with_logits 计算二元交叉熵损失。 reduction="none" 表示不对损失进行平均或求和,返回每个元素的损失值。
loss = F.binary_cross_entropy_with_logits(pred, label, reduction="none")
# 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
# 将预测的 logits 转换为概率值。
pred_prob = pred.sigmoid() # prob from logits
# 计算每个样本的 p_t 值。对于正样本( label=1 ), p_t 是预测概率;对于负样本( label=0 ), p_t 是 1 - 预测概率。
p_t = label * pred_prob + (1 - label) * (1 - pred_prob)
# 计算调节因子,用于减少易分类样本的权重,增加难分类样本的权重。
modulating_factor = (1.0 - p_t) ** gamma
# 将调节因子乘以基础损失,得到加权损失。
loss *= modulating_factor
# 如果 alpha 大于 0,则计算 alpha 因子。alpha 因子用于调整正负样本的权重。
if alpha > 0:
# 计算 alpha 因子。对于正样本, alpha_factor 是 alpha ;对于负样本, alpha_factor 是 1 - alpha 。
alpha_factor = label * alpha + (1 - label) * (1 - alpha)
# 将 alpha 因子乘以加权损失,进一步调整损失值。
loss *= alpha_factor
# 返回最终损失。 loss.mean(1) 对每个样本的损失值求平均, 1 表示沿着第1个维度(通常是特征维度)进行操作。 .sum() 对所有样本的平均损失值求和,得到最终的损失值。
return loss.mean(1).sum()
# 焦点损失通过调节因子和 alpha 因子,减少了易分类样本的权重,增加了难分类样本的权重。这使得模型在训练过程中更加关注难以分类的样本,从而提高了对少数类的分类性能。这种损失函数在目标检测和分类任务中非常有效,能够帮助模型更好地处理类别不平衡问题。
# FocalLoss 类实现了一种用于处理类别不平衡问题的损失函数。通过调节因子和 alpha 因子,它可以减少易分类样本的权重,增加难分类样本的权重,从而提高模型对少数类的分类性能。这种损失函数在目标检测等任务中非常有效,能够帮助模型更好地关注难以分类的样本。
# 这段代码定义了一个名为 BboxLoss 的类,继承自 PyTorch 的 nn.Module ,用于计算边界框(bounding box)的损失。这个类结合了 IoU 损失和分布焦点损失(DFL),用于在目标检测任务中优化模型的边界框预测。
# 定义了一个名为 BboxLoss 的类,继承自 nn.Module 。
class BboxLoss(nn.Module):
# 训练期间计算训练损失的标准类。
"""Criterion class for computing training losses during training."""
# 这段代码是 BboxLoss 类的初始化方法 __init__ 的实现。
# 定义了类的初始化方法。初始化方法用于设置类的属性,并在创建类的实例时被调用。
# 1.self :对当前对象的引用,用于访问类的属性和方法。
# 2.reg_max :一个参数,表示分布焦点损失(DFL)的最大值。这个参数用于定义预测分布的范围。
# 3.use_dfl :一个布尔类型的参数,默认值为 False ,表示是否使用分布焦点损失(DFL)。如果设置为 True ,则在计算损失时会考虑 DFL 损失。
def __init__(self, reg_max, use_dfl=False):
# 使用正则化最大值和 DFL 设置初始化 BboxLoss 模块。
"""Initialize the BboxLoss module with regularization maximum and DFL settings."""
# 调用父类 nn.Module 的初始化方法。这是必要的,因为它确保了 BboxLoss 类正确地继承了 nn.Module 的所有属性和方法。 nn.Module 是 PyTorch 中所有神经网络模块的基类,提供了许多用于构建和管理神经网络模块的功能。
super().__init__()
# 将传入的 reg_max 参数值赋给类的 reg_max 属性。这个属性用于在后续的损失计算中 确定分布的最大值 。
self.reg_max = reg_max
# 将传入的 use_dfl 参数值赋给类的 use_dfl 属性。这个属性用于在损失计算时 决定是否使用分布焦点损失(DFL) 。
self.use_dfl = use_dfl
# 通过这个初始化方法, BboxLoss 类的实例在创建时会根据传入的参数设置其属性。这些属性在后续的损失计算中起着重要作用,特别是当涉及到分布焦点损失(DFL)的计算时。通过设置 use_dfl 参数,用户可以选择是否在损失计算中包含 DFL 损失,从而根据具体的应用需求调整损失函数的行为。
# 这段代码是 BboxLoss 类的 forward 方法的实现,用于计算边界框的损失。这个方法结合了 IoU 损失和分布焦点损失(DFL),用于优化目标检测任务中的边界框预测。
# 定义了前向传播方法 forward ,它接受以下参数 :
# 1.pred_dist :预测的分布值。
# 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损失。
"""IoU loss."""
# 计算权重。target_scores.sum(-1) 对目标分数在最后一个维度上求和,得到每个样本的总分数。 [fg_mask] 使用前景掩码筛选出前景样本的分数。 .unsqueeze(-1) 在最后一个维度上增加一个维度,以便后续的广播操作。
weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
# 计算 IoU 损失。
# bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True) 计算预测边界框和目标边界框之间的 IoU(交并比),使用 CIoU(完整交并比)。
iou = bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True)
# ((1.0 - iou) * weight).sum() 计算加权的 (1 - IoU) 损失,然后对所有样本求和。
# / target_scores_sum 将总损失除以目标分数的总和,得到平均的 IoU 损失。
loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum
# DFL loss
# 如果 use_dfl 为 True ,则计算分布焦点损失(DFL)。
if self.use_dfl:
# 将目标边界框转换为分布值。
# def bbox2dist(anchor_points, bbox, reg_max): -> 用于将边界框转换为相对于锚点的距离表示。 -> 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)
# 调用 _df_loss 方法计算 DFL 损失,并乘以权重。
loss_dfl = self._df_loss(pred_dist[fg_mask].view(-1, self.reg_max + 1), target_ltrb[fg_mask]) * weight
# 计算加权 DFL 损失的总和并除以目标分数的总和。
loss_dfl = loss_dfl.sum() / target_scores_sum
# 如果 use_dfl 为 False ,则将 loss_dfl 设置为 0,并确保其在相同的设备上。
else:
loss_dfl = torch.tensor(0.0).to(pred_dist.device)
# 返回计算得到的 IoU 损失和 DFL 损失。
return loss_iou, loss_dfl
# BboxLoss 类的 forward 方法通过计算 IoU 损失和分布焦点损失(DFL),优化了目标检测任务中的边界框预测。IoU 损失用于衡量预测边界框与目标边界框的重叠程度,而 DFL 损失用于优化边界框的分布表示。通过这种组合,模型可以更好地学习边界框的精确位置和形状,从而提高目标检测的准确性。
# 这段代码定义了 BboxLoss 类中的一个静态方法 _df_loss ,用于计算分布焦点损失(DFL)。
# @staticmethod 表示这是一个静态方法,不需要实例化对象即可调用。
@staticmethod
# 定义了 _df_loss 方法,它接受以下参数 :
# 1.pred_dist :预测的分布值,通常是一个形状为 (N, reg_max + 1) 的张量,其中 N 是样本数量, reg_max 是分布的最大值。
# 2.target :目标值,通常是一个形状为 (N, 4) 的张量,表示每个样本的目标边界框的分布表示。
def _df_loss(pred_dist, target):
# 返回左、右 DFL 损失之和。
# 分布焦点损失 (DFL) 在 Generalized Focal Loss 中提出。
# https://ieeexplore.ieee.org/document/9792391
"""
Return sum of left and right DFL losses.
Distribution Focal Loss (DFL) proposed in Generalized Focal Loss
https://ieeexplore.ieee.org/document/9792391
"""
# 将目标值转换为整数部分,表示目标值的左边界。
tl = target.long() # target left
# 计算目标值的右边界,即整数部分加 1。
tr = tl + 1 # target right
# 计算右边界与目标值之间的差值,作为左边界权重。
wl = tr - target # weight left
# 计算右边界权重,即 1 减去左边界权重。
wr = 1 - wl # weight right
return (
# 计算 预测分布 与 目标左边界 之间的交叉熵损失。 reduction="none" 表示不对损失进行求和或平均,返回每个元素的损失值。
F.cross_entropy(pred_dist, tl.view(-1), reduction="none").view(tl.shape) * wl
# 计算 预测分布 与 目标右边界 之间的交叉熵损失。
# * wl 和 * wr 分别将 左边界 和 右边界 损失乘以相应的 权重 。
# .view(tl.shape) 将损失值的形状调整为目标值的形状,以便进行后续的加权操作。
# + 将加权后的左边界和右边界损失相加。
+ F.cross_entropy(pred_dist, tr.view(-1), reduction="none").view(tl.shape) * wr
# .mean(-1, keepdim=True) 对最后一个维度(通常是特征维度)求平均,并保持该维度,以便返回一个形状为 (N, 1) 的张量,表示每个样本的 DFL 损失。
).mean(-1, keepdim=True)
# _df_loss 方法通过计算预测分布与目标值之间的交叉熵损失,并结合目标值的整数部分和小数部分的权重,来计算分布焦点损失(DFL)。这种方法能够更好地处理边界框的回归问题,尤其是在目标检测任务中,能够提高边界框预测的精度。
# BboxLoss 类结合了 IoU 损失和分布焦点损失(DFL),用于优化目标检测任务中的边界框预测。IoU 损失用于衡量预测边界框与目标边界框的重叠程度,而 DFL 损失用于优化边界框的分布表示。通过这种组合,模型可以更好地学习边界框的精确位置和形状。
# 这段代码定义了一个名为 RotatedBboxLoss 的类,继承自 BboxLoss 类,用于计算旋转边界框的损失。这个类专门用于处理旋转目标检测任务中的损失计算。
# 定义了一个名为 RotatedBboxLoss 的类,继承自 BboxLoss 类。这意味着 RotatedBboxLoss 类继承了 BboxLoss 类的所有属性和方法,并可以对其进行扩展或修改。
class RotatedBboxLoss(BboxLoss):
# 训练期间计算训练损失的标准类。
"""Criterion class for computing training losses during training."""
# 定义了 RotatedBboxLoss 类的初始化方法,接受以下参数 :
# 1.reg_max :表示分布焦点损失(DFL)的最大值。
# 2.use_dfl :一个布尔值,表示是否使用分布焦点损失,默认为 False 。
def __init__(self, reg_max, use_dfl=False):
# 使用正则化最大值和 DFL 设置初始化 BboxLoss 模块。
"""Initialize the BboxLoss module with regularization maximum and DFL settings."""
# 调用父类 BboxLoss 的初始化方法,将 reg_max 和 use_dfl 参数传递给父类。这确保了 RotatedBboxLoss 类正确地继承了 BboxLoss 类的属性和方法。
super().__init__(reg_max, use_dfl)
# 定义了前向传播方法 forward ,它接受以下参数 :
# 1.pred_dist :预测的分布值。
# 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."""
# 计算权重,基于目标分数的总和,并使用前景掩码进行筛选。
weight = target_scores.sum(-1)[fg_mask].unsqueeze(-1)
# 计算预测旋转边界框和目标旋转边界框之间的 IoU(交并比),使用 probiou 函数进行计算。 probiou 是一种用于旋转边界框的 IoU 计算方法。
# def probiou(obb1, obb2, CIoU=False, eps=1e-7):
# -> 用于计算两个旋转边界框( OBBs)之间的概率交并比( ProbIoU)。返回完整交并比(CIoU)。如果 CIoU 为 False ,则直接返回概率交并比 iou 。
# -> return iou - v * alpha # CIoU / return iou
iou = probiou(pred_bboxes[fg_mask], target_bboxes[fg_mask])
# 计算 IoU 损失,通过将 (1 - IoU) 乘以权重并求和,然后除以目标分数的总和。
loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum
# DFL loss
# 如果 use_dfl 为 True ,则计算分布焦点损失(DFL)。
if self.use_dfl:
# 将目标边界框转换为分布值。这里使用 xywh2xyxy 函数将旋转边界框的 (x, y, w, h) 格式转换为 (x1, y1, x2, y2) 格式。
# def bbox2dist(anchor_points, bbox, reg_max): -> 用于将边界框转换为相对于锚点的距离表示。 -> return torch.cat((anchor_points - x1y1, x2y2 - anchor_points), -1).clamp_(0, reg_max - 0.01) # dist (lt, rb)
target_ltrb = bbox2dist(anchor_points, xywh2xyxy(target_bboxes[..., :4]), self.reg_max)
# 调用 _df_loss 方法计算 DFL 损失,并乘以权重。
loss_dfl = self._df_loss(pred_dist[fg_mask].view(-1, self.reg_max + 1), target_ltrb[fg_mask]) * weight
# 计算加权 DFL 损失的总和并除以目标分数的总和。
loss_dfl = loss_dfl.sum() / target_scores_sum
# 如果 use_dfl 为 False ,则将 loss_dfl 设置为 0,并确保其在相同的设备上。
else:
loss_dfl = torch.tensor(0.0).to(pred_dist.device)
# 返回计算得到的 IoU 损失和 DFL 损失。
return loss_iou, loss_dfl
# RotatedBboxLoss 类通过继承 BboxLoss 类并对其进行扩展,实现了旋转边界框的损失计算。它使用 probiou 函数来计算旋转边界框的 IoU 损失,并在需要时计算分布焦点损失(DFL)。这种损失函数适用于旋转目标检测任务,能够帮助模型更好地学习旋转边界框的精确位置和形状,从而提高旋转目标检测的准确性。
# 这段代码定义了一个名为 KeypointLoss 的类,继承自 PyTorch 的 nn.Module ,用于计算关键点检测任务中的损失。这个损失函数基于关键点的欧氏距离和一个缩放因子,通常用于人体姿态估计等任务中。
# 定义了一个名为 KeypointLoss 的类,继承自 nn.Module 。
class KeypointLoss(nn.Module):
# 计算训练损失的标准类别。
"""Criterion class for computing training losses."""
# 义了类的初始化方法,接受一个参数。
# 1.sigmas :表示每个关键点的标准差。
def __init__(self, sigmas) -> None:
# 初始化 KeypointLoss 类。
"""Initialize the KeypointLoss class."""
# 调用父类 nn.Module 的初始化方法,确保 KeypointLoss 类正确地继承了 nn.Module 的所有属性和方法。
super().__init__()
# 将传入的 sigmas 参数值赋给类的 sigmas 属性。
self.sigmas = sigmas
# 定义了前向传播方法 forward ,它接受以下参数 :
# 1.pred_kpts :预测的关键点坐标。
# 2.gt_kpts :真实的关键点坐标。
# 3.kpt_mask :关键点掩码,用于指示哪些关键点是有效的。
# 4.area :关键点所在区域的面积。
def forward(self, pred_kpts, gt_kpts, kpt_mask, area):
# 计算预测和实际关键点的关键点损失因子和欧几里得距离损失。
"""Calculates keypoint loss factor and Euclidean distance loss for predicted and actual keypoints."""
# 计算预测关键点和真实关键点之间的欧氏距离的平方。 pred_kpts[..., 0] 和 gt_kpts[..., 0] 分别表示预测和真实关键点的 x 坐标, pred_kpts[..., 1] 和 gt_kpts[..., 1] 分别表示 y 坐标。
d = (pred_kpts[..., 0] - gt_kpts[..., 0]).pow(2) + (pred_kpts[..., 1] - gt_kpts[..., 1]).pow(2)
# 计算关键点损失因子 kpt_loss_factor ,用于平衡不同样本中有效关键点的数量。 kpt_mask.shape[1] 表示关键点的总数, torch.sum(kpt_mask != 0, dim=1) 计算每个样本中有效关键点的数量, + 1e-9 用于防止除零错误。
kpt_loss_factor = kpt_mask.shape[1] / (torch.sum(kpt_mask != 0, dim=1) + 1e-9)
# e = d / (2 * (area * self.sigmas) ** 2 + 1e-9) # from formula
# 计算损失的缩放因子 e 。这里使用了 COCO 评估中的公式,将欧氏距离 d 除以关键点的标准差 self.sigmas 的平方和区域面积 area 的乘积。
e = d / ((2 * self.sigmas).pow(2) * (area + 1e-9) * 2) # from cocoeval
# (1 - torch.exp(-e)) * kpt_mask 计算每个关键点的损失,使用指数函数对缩放因子 e 进行缩放,并乘以关键点掩码 kpt_mask 以忽略无效关键点。
# kpt_loss_factor.view(-1, 1) * ... 将关键点损失因子与每个关键点的损失相乘,以平衡不同样本中有效关键点的数量。
# .mean() 对所有样本和关键点的损失求平均,得到最终的损失值。
return (kpt_loss_factor.view(-1, 1) * ((1 - torch.exp(-e)) * kpt_mask)).mean()
# KeypointLoss 类通过计算关键点的欧氏距离和缩放因子,提供了一种用于关键点检测任务的损失函数。这种损失函数能够有效地衡量预测关键点与真实关键点之间的差异,并通过关键点掩码和损失因子来处理不同样本中有效关键点数量的不平衡问题。这种损失函数在人体姿态估计等任务中非常有用,能够帮助模型更好地学习关键点的精确位置,从而提高关键点检测的准确性。
# 这段代码定义了一个名为 v8DetectionLoss 的类,用于计算目标检测任务中的训练损失。这个类结合了多种损失类型,包括边界框损失、分类损失和分布焦点损失(DFL),以优化模型的检测性能。
# 定义了一个名为 v8DetectionLoss 的类。
class v8DetectionLoss:
# 计算训练损失的标准类别。
"""Criterion class for computing training losses."""
# 这段代码是 v8DetectionLoss 类的初始化方法 __init__ 的实现。它用于设置损失函数所需的各个组件和参数。
# 定义了类的初始化方法,接受两个参数。
# 1.model :检测模型,用于获取模型的参数和配置。
# 2.tal_topk :任务对齐分配器的参数,表示在分配正样本时考虑的 top-k 个候选框,默认值为 10。
# 注释 # model must be de-paralleled 表示传入的模型应该是非并行化的,即不是通过 nn.DataParallel 或 nn.parallel.DistributedDataParallel 包装过的模型。
def __init__(self, model, tal_topk=10): # model must be de-paralleled
# 使用模型初始化 v8DetectionLoss,定义与模型相关的属性和 BCE 损失函数。
"""Initializes v8DetectionLoss with the model, defining model-related properties and BCE loss function."""
# 获取模型参数所在的设备(如 CPU 或 GPU)。
device = next(model.parameters()).device # get model device
# 获取模型的超参数配置。
h = model.args # hyperparameters
# 获取模型的最后一个模块,通常是一个检测模块(如 Detect() ),它包含了检测任务相关的配置和属性。
m = model.model[-1] # Detect() module
# 初始化损失函数和属性。
# 初始化二元交叉熵损失函数,用于 分类损失 的计算。 reduction="none" 表示不对损失进行求和或平均,返回每个元素的损失值。
self.bce = nn.BCEWithLogitsLoss(reduction="none")
# 将模型的 超参数配置 赋给类的 hyp 属性。
self.hyp = h
# 获取检测模块的 步长 (stride),用于计算特征图与输入图像的尺寸关系。
self.stride = m.stride # model strides
# 获取检测模块的 类别数 。
self.nc = m.nc # number of classes
# 获取检测模块的 输出维度 。
self.no = m.no
# 获取检测模块的 最大回归值 。
self.reg_max = m.reg_max
# 将模型参数所在的设备赋给类的 device 属性。
self.device = device
# 根据最大回归值 reg_max 判断是否使用分布焦点损失(DFL)。如果 reg_max 大于 1,则使用 DFL。
self.use_dfl = m.reg_max > 1
# 初始化任务对齐分配器,用于分配正负样本。传入参数包括 topk 、类别数 num_classes 、以及用于计算对齐度量的 alpha 和 beta 。
self.assigner = TaskAlignedAssigner(topk=tal_topk, num_classes=self.nc, alpha=0.5, beta=6.0)
# 初始化边界框损失函数,并将其移动到模型参数所在的设备上。
self.bbox_loss = BboxLoss(m.reg_max - 1, use_dfl=self.use_dfl).to(device)
# 创建一个投影张量,用于 DFL 损失的计算。该张量包含从 0 到 reg_max - 1 的整数,数据类型为 float ,并位于模型参数所在的设备上。
self.proj = torch.arange(m.reg_max, dtype=torch.float, device=device)
# 通过这个初始化方法, v8DetectionLoss 类的实例在创建时会根据传入的模型和参数设置其属性和组件。这些属性和组件在后续的损失计算中起着重要作用,特别是当涉及到分类损失、边界框回归损失和分布焦点损失(DFL)的计算时。通过设置这些属性和组件,损失函数能够根据模型的配置和任务需求进行相应的计算和优化。
# 这段代码定义了 v8DetectionLoss 类的 preprocess 方法,用于对目标数据进行预处理。这个方法将原始的目标数据转换为适合后续损失计算的格式。
# 定义了 preprocess 方法,它接受以下参数 :
# 1.targets :目标数据,通常是一个形状为 (N, 6) 的张量,其中 N 是目标的数量,每一行包含 图像索引 、 类别标签 和边界框坐标( x, y, w, h )。
# 2.batch_size :批次大小。
# 3.scale_tensor :用于调整边界框坐标的缩放因子。
def preprocess(self, targets, batch_size, scale_tensor):
# 预处理目标计数并与输入批次大小匹配以输出张量。
"""Preprocesses the target counts and matches with the input batch size to output a tensor."""
# 如果目标数据为空(即 targets.shape[0] == 0 ),则返回一个形状为 (batch_size, 0, 5) 的全零张量。这表示当前批次中没有目标。
if targets.shape[0] == 0:
out = torch.zeros(batch_size, 0, 5, device=self.device)
# 处理非空目标。
else:
# 获取目标数据中的图像索引。
i = targets[:, 0] # image index
# 计算每个图像索引的出现次数,即每个图像中的目标数量。
_, counts = i.unique(return_counts=True)
# 将目标数量转换为 int32 类型。
counts = counts.to(dtype=torch.int32)
# 创建一个输出张量,形状为 (batch_size, max_targets, 5) ,其中 max_targets 是每个图像中目标的最大数量。这个张量用于存储预处理后的目标数据。
out = torch.zeros(batch_size, counts.max(), 5, device=self.device)
# 分配目标到批次。
# 遍历每个批次中的图像索引 j 。
for j in range(batch_size):
# 找出目标数据中属于当前图像索引 j 的目标。
matches = i == j
# 计算当前图像中的目标数量。
n = matches.sum()
# 如果 n 大于 0,则将这些目标的数据(不包括图像索引)复制到输出张量 out 中对应的位置。
if n:
out[j, :n] = targets[matches, 1:]
# 将目标的边界框坐标乘以缩放因子 scale_tensor ,以调整其大小。 xywh2xyxy(...) 将边界框的格式从 (x, y, w, h) 转换为 (x1, y1, x2, y2) ,其中 (x1, y1) 是左上角坐标, (x2, y2) 是右下角坐标。
out[..., 1:5] = xywh2xyxy(out[..., 1:5].mul_(scale_tensor))
# 返回预处理后的目标数据 out ,它包含了每个批次中每个图像的目标类别和边界框坐标。
return out
# preprocess 方法通过对目标数据进行预处理,将其转换为适合后续损失计算的格式。它能够处理空目标的情况,并将非空目标分配到相应的批次中,同时将边界框格式从 (x, y, w, h) 转换为 (x1, y1, x2, y2) 。这种预处理方式为损失函数的计算提供了便利,使得损失函数能够更有效地处理目标检测任务中的目标数据。
# 这段代码定义了 v8DetectionLoss 类的 bbox_decode 方法,用于将预测的分布值解码为边界框坐标。这个方法特别适用于使用分布焦点损失(DFL)的场景,能够将预测的分布值转换为具体的边界框位置。
# 定义了 bbox_decode 方法,它接受以下参数 :
# 1.anchor_points :锚点坐标,通常是一个形状为 (N, 2) 的张量,表示每个锚点的中心坐标 (x, y) 。
# 2.pred_dist :预测的分布值,通常是一个形状为 (batch_size, num_anchors, num_channels) 的张量,表示每个锚点的预测分布。
def bbox_decode(self, anchor_points, pred_dist):
# 从锚点和分布解码预测的对象边界框坐标。
"""Decode predicted object bounding box coordinates from anchor points and distribution."""
# 判断是否使用分布焦点损失(DFL)。如果使用 DFL,则进行以下操作
if self.use_dfl:
# 获取预测分布值的形状,分别表示 批次大小 b 、 锚点数量 a 和 通道数 c 。
b, a, c = pred_dist.shape # batch, anchors, channels
# 这行代码是 bbox_decode 方法中用于将预测的分布值转换为实际距离值的关键步骤。
# pred_dist.view(b, a, 4, c // 4) 将预测的分布值 pred_dist 重塑为一个新的形状 (b, a, 4, c // 4) 。
# b 是批次大小。 a 是锚点数量。 4 表示边界框的四个方向(左、右、上、下)。 c // 4 表示每个方向上的离散距离值的数量。 c 是原始预测分布值的通道数, c // 4 表示每个方向上的分布值数量。
# .softmax(3) 对重塑后的张量在最后一个维度(即每个方向上的分布值)应用 softmax 函数。这一步将每个方向上的分布值转换为概率分布,使得每个方向上的所有值的和为 1。
# 例如,对于某个方向的分布值 [0.5, 0.3, 0.2] ,softmax 函数会将其转换为概率分布,如 [0.5744, 0.2447, 0.1809] 。
# .matmul(self.proj.type(pred_dist.dtype)) 将 softmax 处理后的概率分布与投影张量 self.proj 进行矩阵乘法。
# self.proj.type(pred_dist.dtype) 确保投影张量 self.proj 的数据类型与 pred_dist 一致,以避免类型不匹配导致的错误。
# 这一步将每个方向的概率分布与对应的离散距离值相乘,并求和,得到每个方向上的加权平均距离值。
# 例如,假设某个方向的概率分布为 [0.5744, 0.2447, 0.1809] ,对应的离散距离值为 [0, 1, 2] ,则加权平均距离值为:0×0.5744+1×0.2447+2×0.1809=0.6065 。
# 这行代码通过重塑张量、应用 softmax 函数和矩阵乘法,将预测的分布值转换为实际的距离值。这种方法使得模型能够从预测的分布中恢复出具体的边界框位置,提高了边界框预测的精度和稳定性。通过这种方式,模型能够更好地处理边界框的回归问题,尤其是在使用分布焦点损失(DFL)的场景中。
# 在使用分布焦点损失(DFL)的上下文中,投影张量 self.proj 的作用是将预测的分布值转换为实际的距离值。具体来说, self.proj 是一个从 0 到 reg_max - 1 的整数序列,用于将每个方向的预测分布转换为具体的距离值。以下是其具体作用的详细解释 :
# 分布焦点损失(DFL)的背景 :
# 分布焦点损失(DFL) 是一种用于边界框回归的损失函数,它通过预测边界框的分布来提高回归精度。DFL 的核心思想是将边界框的每个方向(如左、右、上、下)的回归问题转化为一个分类问题,即预测该方向上的距离分布。
# 在 DFL 中,每个方向的距离被离散化为多个可能的值,模型的任务是预测这些值的分布概率。
# 投影张量 self.proj 的作用 :
# 定义 : self.proj 是一个形状为 (reg_max,) 的张量,包含了从 0 到 reg_max - 1 的整数序列。 reg_max 是离散化的最大值,表示边界框在每个方向上可以取的最大距离值。
# 转换过程 :
# 在解码过程中,模型首先将每个方向的预测分布值通过 softmax 函数转换为概率分布。
# 然后,通过与 self.proj 进行矩阵乘法,将这些概率分布转换为加权平均的距离值。具体来说,每个方向的预测距离是所有可能距离值的加权和,权重为该距离值对应的概率。
# 这种方法允许模型在预测边界框时考虑多个可能的距离值,从而提高了预测的精度和鲁棒性。
# 示例 :
# 假设 reg_max = 4 ,则 self.proj 是 [0, 1, 2, 3] 。如果某个方向的预测分布概率为 [0.1, 0.6, 0.2, 0.1] ,则通过矩阵乘法计算得到的预测距离为:0×0.1 + 1×0.6 + 2×0.2 + 3×0.1 = 1.1 这个值表示该方向上的预测距离。
# 总结 :
# 投影张量 self.proj 在 DFL 中起到了将预测的分布值转换为实际距离值的关键作用。通过这种方式,模型能够更好地处理边界框的回归问题,提高预测的精度和稳定性。
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) ,即左上角和右下角坐标。
# def dist2bbox(distance, anchor_points, xywh=True, dim=-1): -> 用于将距离(ltrb,即左、上、右、下)转换为边界框(bbox)的坐标。 -> return torch.cat((c_xy, wh), dim) # xywh bbox / return torch.cat((x1y1, x2y2), dim) # xyxy bbox
return dist2bbox(pred_dist, anchor_points, xywh=False)
# bbox_decode 方法通过将预测的分布值解码为边界框坐标,使得模型能够从预测的分布中恢复出具体的边界框位置。这种方法特别适用于使用分布焦点损失(DFL)的场景,能够提高边界框预测的精度和稳定性。通过将分布值转换为实际的距离,并结合锚点坐标,模型能够更准确地定位目标的位置。
# 这段代码是 v8DetectionLoss 类的 __call__ 方法的实现,用于计算目标检测任务中的训练损失。这个方法结合了多种损失函数,包括边界框回归损失、分类损失和分布焦点损失(DFL),以优化模型的检测性能。
# 定义了 __call__ 方法,它接受两个参数。
# 1.preds :模型的预测输出,可能是一个张量或一个包含多个张量的元组。
# 2.batch :批次数据,包含图像索引、类别标签和边界框坐标等信息。
def __call__(self, preds, batch):
# 计算 box、cls 和 dfl 的损失之和乘以 batch size。
"""Calculate the sum of the loss for box, cls and dfl multiplied by batch size."""
# 这段代码是 v8DetectionLoss 类的 __call__ 方法的初始部分,用于初始化损失张量并处理模型的预测输出。
# 初始化一个形状为 (3,) 的张量 loss ,用于存储三种类型的损失 : loss[0] 边界框回归损失(box loss)。 loss[1] 分类损失(classification loss)。 loss[2] 分布焦点损失(distribution focal loss,DFL)。
# device=self.device 确保损失张量位于与模型相同的设备上(如 CPU 或 GPU),以便进行后续的计算。
loss = torch.zeros(3, device=self.device) # box, cls, dfl
# 处理模型的预测输出 preds 。如果 preds 是一个元组(tuple),则取元组的第二个元素作为特征图 feats 。这通常发生在模型输出包含多个部分的情况下,例如,除了特征图外,还可能包含其他辅助信息。 如果 preds 不是元组,则直接将其视为特征图 feats 。
feats = preds[1] if isinstance(preds, tuple) else preds
# 将特征图 feats 中的预测输出分割为 预测分布值 pred_distri 和预测分数 pred_scores 。
# torch.cat([xi.view(feats[0].shape[0], self.no, -1) for xi in feats], 2) : xi.view(feats[0].shape[0], self.no, -1) 将每个特征图 xi 重塑为形状 (batch_size, num_outputs, -1) ,其中 batch_size 是批次大小, num_outputs 是每个特征图的输出维度(由 self.no 指定), -1 表示自动计算剩余维度。
# torch.cat(..., 2) 将所有特征图沿着第三个维度(即特征维度)拼接起来,形成一个完整的预测输出张量。
# split((self.reg_max * 4, self.nc), 1) :将拼接后的预测输出张量分割为两部分 :
# pred_distri :预测分布值,形状为 (batch_size, num_anchors, self.reg_max * 4) ,包含边界框的分布信息。
# pred_scores :预测分数,形状为 (batch_size, num_anchors, self.nc) ,包含每个锚点的类别分数。
# self.reg_max * 4 和 self.nc 分别指定了分割的位置, self.reg_max 是分布的最大值, 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
)
# 这段代码通过初始化损失张量和处理模型的预测输出,为后续的损失计算奠定了基础。它将模型的预测输出分割为边界框的分布值和类别分数,以便分别计算边界框回归损失、分类损失和分布焦点损失(DFL)。这种处理方式使得损失函数能够全面地评估模型在目标检测任务中的性能。
# 这段代码继续处理模型的预测输出,并准备一些必要的参数和数据结构,以便进行后续的损失计算。
# 调整 pred_scores 和 pred_distri 的维度顺序 :
# permute(0, 2, 1) 将张量的维度顺序从 (batch_size, num_anchors, num_classes) 或 (batch_size, num_anchors, num_channels) 调整为 (batch_size, num_classes, num_anchors) 或 (batch_size, num_channels, num_anchors) 。这种调整通常是为了使张量的布局更适合后续的计算,例如矩阵乘法或损失函数的计算。
# contiguous() 返回一个连续的内存布局的张量副本。在某些操作(如 view 或 reshape )之前,调用 contiguous() 可以确保张量在内存中是连续存储的,从而避免潜在的错误或性能问题。
pred_scores = pred_scores.permute(0, 2, 1).contiguous()
pred_distri = pred_distri.permute(0, 2, 1).contiguous()
# 获取预测分数张量 pred_scores 的数据类型。这通常是为了确保后续计算中使用相同的数据类型,以避免类型不匹配导致的错误。
dtype = pred_scores.dtype
# 获取批次大小,即 pred_scores 的第一个维度的大小。批次大小表示当前批次中图像的数量。
batch_size = pred_scores.shape[0]
# 创建一个张量,包含特征图 feats[0] 的空间维度大小(即高度和宽度),并将其放置在与模型相同的设备上,并使用与 pred_scores 相同的数据类型。
# * self.stride[0] 将特征图的空间维度大小乘以模型的步幅 self.stride[0] ,以计算原始输入图像的尺寸(高度和宽度)。步幅表示特征图相对于输入图像的缩小比例,因此通过乘以步幅可以恢复到输入图像的实际尺寸。
imgsz = torch.tensor(feats[0].shape[2:], device=self.device, dtype=dtype) * self.stride[0] # image size (h,w)
# 调用 make_anchors 函数,用于生成 锚点坐标 anchor_points 和 步幅张量 stride_tensor 。
# feats 是特征图的列表或张量,用于计算锚点的位置。 self.stride 是模型的步幅列表,表示特征图相对于输入图像的缩小比例。 0.5 是一个偏移量,通常用于调整锚点的位置,使其更接近特征图的中心。
# 锚点坐标 anchor_points 表示在输入图像上用于检测目标的预定义位置, 步幅张量 stride_tensor 表示每个锚点相对于特征图的步幅。
# def make_anchors(feats, strides, grid_cell_offset=0.5): -> 用于生成目标检测模型中使用的锚点(anchors)。返回连接后的 锚点坐标张量 和 步长张量 。 -> return torch.cat(anchor_points), torch.cat(stride_tensor)
anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5)
# 这段代码通过调整预测分数和分布值的张量维度顺序,获取必要的参数(如数据类型、批次大小和图像尺寸),并生成锚点坐标和步幅张量,为后续的损失计算做好了准备。这些步骤确保了数据结构的一致性和计算的准确性,使得损失函数能够有效地评估模型在目标检测任务中的性能。
# 这段代码处理目标数据,将其转换为适合损失计算的格式,并提取出类别标签和边界框坐标。
# Targets
# batch["batch_idx"].view(-1, 1) 将批次索引 batch_idx 转换为形状为 (-1, 1) 的张量,其中 -1 表示自动计算维度以匹配数据的总数。
# batch["cls"].view(-1, 1) 将类别标签 cls 转换为形状为 (-1, 1) 的张量。
# torch.cat((..., ...), 1) 将 批次索引 、 类别标签 和 边界框坐标 batch["bboxes"] 沿着第二个维度(即列)拼接起来,形成一个形状为 (N, 6) 的张量 targets ,其中 N 是目标的数量。每一行包含一个目标的信息,格式为 [batch_idx, cls, x, y, w, h] 。
targets = torch.cat((batch["batch_idx"].view(-1, 1), batch["cls"].view(-1, 1), batch["bboxes"]), 1)
# targets.to(self.device) 将目标数据移动到与模型相同的设备上(如 CPU 或 GPU)。
# self.preprocess(...) 调用 preprocess 方法对目标数据进行预处理。预处理过程包括 :
# 将目标数据分配到不同的批次中。
# 将边界框坐标从 (x, y, w, h) 格式转换为 (x1, y1, x2, y2) 格式。
# 使用 scale_tensor=imgsz[[1, 0, 1, 0]] 对边界框坐标进行缩放,以匹配输入图像的尺寸。 imgsz[[1, 0, 1, 0]] 表示将图像尺寸的顺序调整为 (w, h, w, h) ,以适应边界框坐标的转换。
targets = self.preprocess(targets.to(self.device), batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
# 将预处理后的目标数据 targets 沿着第三个维度(即特征维度)分割为两部分 :
# gt_labels :类别标签,形状为 (batch_size, max_targets, 1) 。
# gt_bboxes :边界框坐标,形状为 (batch_size, max_targets, 4) ,格式为 (x1, y1, x2, y2) 。
gt_labels, gt_bboxes = targets.split((1, 4), 2) # cls, xyxy
# gt_bboxes.sum(2, keepdim=True) 对每个边界框的坐标求和,并保持维度不变,结果形状为 (batch_size, max_targets, 1) 。
# .gt_(0) 生成一个掩码张量 mask_gt ,表示哪些边界框是有效的(即坐标和大于 0)。这个掩码用于后续的损失计算,以忽略无效的目标。
mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0)
# 这段代码通过拼接、预处理和分割目标数据,提取出类别标签和边界框坐标,并生成目标掩码。这些步骤为后续的损失计算提供了必要的输入数据,使得损失函数能够准确地评估模型在目标检测任务中的性能。通过这种方式,模型能够更好地学习目标的类别和位置信息。
# 这段代码处理预测的边界框和目标分配,是目标检测损失计算中的关键步骤。
# Pboxes
# 解码预测边界框。
# self.bbox_decode(anchor_points, pred_distri) 调用 bbox_decode 方法将预测的分布值 pred_distri 解码为边界框坐标 pred_bboxes 。
# anchor_points 是生成的锚点坐标,用于帮助解码边界框。
# 解码后的边界框坐标 pred_bboxes 的格式为 (x1, y1, x2, y2) ,形状为 (batch_size, num_anchors, 4) ,表示每个锚点对应的预测边界框。
pred_bboxes = self.bbox_decode(anchor_points, pred_distri) # xyxy, (b, h*w, 4)
# 目标分配。
# target_bboxes 是匹配到的真实边界框坐标。
# target_scores 是匹配到的真实目标分数。
# fg_mask 是前景掩码,表示哪些预测是正样本。
# self.assigner(...) 调用任务对齐分配器 self.assigner 进行目标分配。任务对齐分配器的作用是将预测的边界框与真实的目标进行匹配,确定哪些预测是正样本(与真实目标匹配的预测)。
# def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt):
# -> 将预测的边界框和分类分数与真实标注进行匹配,并计算目标标签、目标边界框、目标分数等。返回 目标标签 、 目标边界框 、 目标分数 、 前景掩码 和 目标真实标注索引 。
# -> return target_labels, target_bboxes, target_scores, fg_mask.bool(), target_gt_idx
_, target_bboxes, target_scores, fg_mask, _ = self.assigner(
# 将预测分数通过 sigmoid 函数转换为概率值,并从计算图中分离出来,以避免梯度传播。
pred_scores.detach().sigmoid(),
# 将解码后的边界框坐标乘以步幅张量 stride_tensor ,并转换为与真实边界框 gt_bboxes 相同的数据类型。这一步是为了将边界框坐标从特征图尺度转换回输入图像尺度。
(pred_bboxes.detach() * stride_tensor).type(gt_bboxes.dtype),
# 将锚点坐标乘以步幅张量,以将锚点坐标转换到输入图像尺度。
anchor_points * stride_tensor,
# gt_labels 、 gt_bboxes 和 mask_gt 分别是 真实的目标类别标签 、 边界框坐标 和 目标掩码 。
gt_labels,
gt_bboxes,
mask_gt,
)
# 计算目标分数总和。
# target_scores.sum() 计算所有匹配到的真实目标分数的总和。
# max(..., 1) 确保目标分数总和至少为 1,以避免后续计算中出现除零错误。这在没有匹配到任何目标的情况下特别有用。
target_scores_sum = max(target_scores.sum(), 1)
# 这段代码通过解码预测的边界框和进行目标分配,为后续的损失计算提供了必要的数据。解码后的边界框用于与真实目标进行匹配,目标分配器确定了哪些预测是正样本,并计算了匹配到的真实目标分数的总和。这些步骤确保了损失函数能够准确地评估模型在目标检测任务中的性能,特别是在分类损失和边界框回归损失的计算中。
# 这段代码计算目标检测任务中的分类损失和边界框损失,并根据超参数对损失进行缩放,最终返回总损失和各个损失的值。
# Cls loss 计算分类损失。
# 注释掉的代码 self.varifocal_loss(...) 表示另一种计算分类损失的方法,即使用变焦损失(Varifocal Loss,VFL)。VFL 是一种用于目标检测的分类损失函数,能够更好地处理类别不平衡问题。
# loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum # VFL way
# 这里使用二元交叉熵损失(BCE)来计算分类损失。
# self.bce(pred_scores, target_scores.to(dtype)) 计算 预测分数 pred_scores 和 目标分数 target_scores 之间的二元交叉熵损失。 target_scores.to(dtype) 确保目标分数的数据类型与预测分数一致。
# .sum() 对所有样本的损失值求和。
# / target_scores_sum 将总损失除以目标分数的总和 target_scores_sum ,以获得平均损失。
loss[1] = self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum # BCE
# Bbox loss 计算边界框损失。
# 如果有正样本(即 fg_mask.sum() 大于 0),则计算边界框损失。
if fg_mask.sum():
# 将目标边界框坐标除以步幅张量 stride_tensor ,以将坐标从输入图像尺度转换回特征图尺度。
target_bboxes /= stride_tensor
# self.bbox_loss(...) 调用边界框损失函数 self.bbox_loss ,计算边界框回归损失 loss[0] 和分布焦点损失(DFL) loss[2] 。传入的参数包括预测的分布值、解码后的边界框坐标、锚点坐标、目标边界框坐标、目标分数、目标分数总和和前景掩码。
# def forward(self, pred_dist, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask):
# -> 用于计算边界框的损失。这个方法结合了 IoU 损失和分布焦点损失(DFL),用于优化目标检测任务中的边界框预测。返回计算得到的 IoU 损失和 DFL 损失。
# -> return loss_iou, loss_dfl
loss[0], loss[2] = self.bbox_loss(
pred_distri, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask
)
# 根据超参数对损失进行缩放。
# 将 边界框回归损失 乘以超参数 self.hyp.box ,以调整损失的权重。
loss[0] *= self.hyp.box # box gain
# 将 分类损失 乘以超参数 self.hyp.cls ,以调整损失的权重。
loss[1] *= self.hyp.cls # cls gain
# 将 分布焦点损失 乘以超参数 self.hyp.dfl ,以调整损失的权重。
loss[2] *= self.hyp.dfl # dfl gain
# 返回总损失和各个损失。 loss.sum() * batch_size 计算总损失,并乘以批次大小 batch_size ,以获得整个批次的损失。 loss.detach() 返回各个损失的值,并从计算图中分离出来,以便进行后续的分析和记录。
return loss.sum() * batch_size, loss.detach() # loss(box, cls, dfl)
# 这段代码通过计算分类损失和边界框损失,并根据超参数对损失进行缩放,提供了目标检测任务中的综合损失函数。这种损失函数能够有效地优化模型的检测性能,使模型能够更好地学习目标的类别和位置。通过这种方式,模型在训练过程中能够获得准确的损失反馈,从而提高目标检测的准确性。
# __call__ 方法通过结合边界框回归损失、分类损失和分布焦点损失(DFL),提供了一种用于目标检测任务的综合损失函数。这种损失函数能够有效地优化模型的检测性能,使模型能够更好地学习目标的类别和位置。通过任务对齐分配器和边界框解码方法,该方法能够处理复杂的检测任务,并在训练过程中提供准确的损失反馈。
# v8DetectionLoss 类通过结合分类损失、边界框回归损失和分布焦点损失(DFL),提供了一种用于目标检测任务的综合损失函数。这种损失函数能够有效地优化模型的检测性能,使模型能够更好地学习目标的类别和位置。通过任务对齐分配器和边界框解码方法,该类能够处理复杂的检测任务,并在训练过程中提供准确的损失反馈。
# 这段代码定义了一个名为 v8SegmentationLoss 的类,它是 v8DetectionLoss 的子类,用于计算分割任务中的训练损失。
class v8SegmentationLoss(v8DetectionLoss):
# 计算训练损失的标准类。
"""Criterion class for computing training losses."""
# 这段代码是 v8SegmentationLoss 类的初始化方法 __init__ 的定义。这个方法在创建类的实例时被调用,用于初始化类的属性。
# 这是 v8SegmentationLoss 类的初始化方法,它接收一个参数。
# 1.model :这个参数是一个模型对象,且必须是非并行化的(即不是通过数据并行或模型并行处理过的模型)。这是因为并行化后的模型在某些属性和方法上可能与非并行化模型有所不同,直接使用可能会导致错误。
def __init__(self, model): # model must be de-paralleled
# 初始化 v8SegmentationLoss 类,以去并行模型作为参数。
"""Initializes the v8SegmentationLoss class, taking a de-paralleled model as argument."""
# 调用了父类 v8DetectionLoss 的初始化方法,并将 model 参数传递给它。这是Python中继承机制的一部分,用于确保子类能够正确地初始化其父类的部分,从而继承父类的属性和方法。
super().__init__(model)
# 从传入的 model 对象中获取 args 属性下的 overlap_mask 值,并将其赋值给当前类的 self.overlap 属性。 overlap_mask 是一个布尔值,用于指示在计算分割损失时是否考虑掩码之间的重叠。这个属性在后续计算分割损失时会被用到。
self.overlap = model.args.overlap_mask
# __init__ 方法是 v8SegmentationLoss 类的构造函数,用于初始化类的实例。 它首先调用父类 v8DetectionLoss 的构造函数,以确保继承父类的初始化逻辑。 然后,它从传入的模型对象中提取 overlap_mask 参数,并将其存储为类的属性 self.overlap ,以便在后续的损失计算中使用。 这个初始化过程确保了 v8SegmentationLoss 类能够正确地设置和使用其父类的属性和方法,同时添加了特定于分割任务的配置。
# 这段代码定义了 v8SegmentationLoss 类的 __call__ 方法,该方法用于计算分割任务的总损失。
# 定义了 __call__ 方法,使得类的实例可以像函数一样被调用。接收两个参数。
# 1.preds :模型的预测值。
# 2.batch :当前批次的数据。
def __call__(self, preds, batch):
# 计算并返回 YOLO 模型的损失。
"""Calculate and return the loss for the YOLO model."""
# 这段代码是 v8SegmentationLoss 类的 __call__ 方法的一部分,主要负责初始化损失张量和处理模型的预测值。
# 初始化一个包含4个元素的零张量,存储在指定的设备上(如CPU或GPU)。这个张量用于存储不同部分的损失 : 边界框损失(box) 、 分类损失(cls) 、 分布焦点损失(dfl) 和 分割损失(seg) 。初始值为零,后续会根据计算结果进行更新。
loss = torch.zeros(4, device=self.device) # box, cls, dfl
# 从模型的预测值 preds 中提取 特征 ( feats )、 预测掩码 ( pred_masks )和 原型 ( proto )。如果 preds 是一个长度为3的元组,直接解包赋值;否则,假设 preds 是一个包含多个元素的元组,取其第二个元素(索引为1)进行解包。这种处理方式是为了适应不同情况下预测值的组织形式,确保能够正确提取所需的预测信息。
feats, pred_masks, proto = preds if len(preds) == 3 else preds[1]
# 从原型张量 proto 的形状中提取 批量大小 ( batch_size )、 掩码高度 ( mask_h )和 掩码宽度 ( mask_w )。 _ 用于忽略原型张量的第二个维度(通常是掩码的数量),因为主要关注批量大小和掩码的空间尺寸。这些尺寸信息在后续的计算中会用到,例如对掩码进行缩放或裁剪等操作。
batch_size, _, mask_h, mask_w = proto.shape # batch size, number of masks, mask height, mask width
# 首先,对特征张量 feats 中的每个元素 xi 进行视图变换和拼接。 xi.view(feats[0].shape[0], self.no, -1) 将每个特征张量重新排列为形状为 (batch_size, self.no, -1) 的张量,其中 self.no 是每个锚点的输出数量, -1 表示自动计算该维度的大小以保持总元素数量不变。然后,使用 torch.cat 将这些重新排列的特征张量沿着第二个维度(索引为2)进行拼接,形成一个更大的张量。
# 接下来,使用 split 方法将拼接后的张量分割为两个部分 : 预测分布 ( pred_distri )和 预测分数 ( pred_scores )。分割的依据是 (self.reg_max * 4, self.nc) ,其中 self.reg_max 是回归的最大值, self.nc 是类别数量。
# split 方法会沿着第一个维度(索引为1)进行分割,将前 self.reg_max * 4 个元素分配给 pred_distri ,剩余的 self.nc 个元素分配给 pred_scores 。这样就得到了模型对边界框分布和类别分数的预测值,为后续的损失计算提供了基础.
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
)
# 这段代码主要负责初始化损失张量和处理模型的预测值。通过提取特征、预测掩码和原型,获取原型的尺寸信息,并对预测分布和分数进行处理,为后续的损失计算奠定了基础。这些操作确保了损失计算能够正确地针对模型的预测结果进行,从而有效地评估模型在分割任务中的性能。
# 这段代码继续处理模型的预测值,并进行一些准备工作,以便后续计算损失。
# B, grids, ..
# 对 预测分数 ( pred_scores )、 预测分布 ( pred_distri ) 和 预测掩码 ( pred_masks )进行维度排列。
# permute(0, 2, 1) 将张量的维度从 (batch_size, num_predictions, num_classes) 或 (batch_size, num_predictions, num_prototypes) 变为 (batch_size, num_classes, num_predictions) 或 (batch_size, num_prototypes, num_predictions) 。
# contiguous() 方法确保张量在内存中是连续的,这对于后续的计算和操作是有益的。
pred_scores = pred_scores.permute(0, 2, 1).contiguous()
pred_distri = pred_distri.permute(0, 2, 1).contiguous()
pred_masks = pred_masks.permute(0, 2, 1).contiguous()
# 获取预测分数的 数据类型 ( dtype ),这通常是为了确保后续计算中数据类型的统一。
dtype = pred_scores.dtype
# 计算图像的尺寸(高度和宽度)。 feats[0].shape[2:] 获取特征图的高度和宽度,然后乘以步长 self.stride[0] ,得到原始图像的尺寸。结果存储在 imgsz 中,数据类型与预测分数相同,确保计算的一致性。
imgsz = torch.tensor(feats[0].shape[2:], device=self.device, dtype=dtype) * self.stride[0] # image size (h,w)
# 调用 make_anchors 函数生成 锚点 ( anchor_points )和 步长张量 ( stride_tensor )。 feats 是特征图列表, self.stride 是每个特征图的步长, 0.5 是锚点的偏移量。生成的锚点和步长张量将用于后续的边界框解码和损失计算。
# def make_anchors(feats, strides, grid_cell_offset=0.5): -> 用于生成目标检测模型中使用的锚点(anchors)。返回连接后的 锚点坐标张量 和 步长张量 。 -> return torch.cat(anchor_points), torch.cat(stride_tensor)
anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5)
# 这段代码对预测分数、预测分布和预测掩码进行维度排列,确保它们的维度顺序适合后续的计算。获取预测分数的数据类型,以便在后续计算中保持数据类型的一致性。计算图像的尺寸,这在后续的边界框解码和损失计算中非常有用。生成锚点和步长张量,这些张量将用于解码预测的边界框,确保边界框的计算是基于正确的锚点和步长。这些准备工作为后续的损失计算提供了必要的基础,确保了计算的准确性和效率。
# 这段代码负责处理训练批次中的目标数据,包括批量索引、类别标签和边界框坐标。它将这些目标数据进行整合和预处理,以便后续的损失计算。
# Targets
try:
# 从批次数据 batch 中获取批量索引 batch_idx ,并将其重塑为形状为 (-1, 1) 的张量。这一步是为了将批量索引与类别标签和边界框坐标进行拼接。
batch_idx = batch["batch_idx"].view(-1, 1)
# 使用 torch.cat 将 批量索引 、 类别标签 和 边界框坐标 沿着第二个维度(索引为1)进行拼接,形成一个统一的目标张量 targets 。类别标签 batch["cls"] 也被重塑为形状为 (-1, 1) 的张量,以便与批量索引和边界框坐标对齐。
targets = torch.cat((batch_idx, batch["cls"].view(-1, 1), batch["bboxes"]), 1)
# 将目标张量 targets 移动到指定的设备(如CPU或GPU),并调用 self.preprocess 方法进行预处理。预处理过程包括对边界框坐标进行缩放等操作,以适应模型的输入要求。 imgsz[[1, 0, 1, 0]] 表示图像尺寸的 xyxy 格式(即宽度、高度、宽度、高度)。
targets = self.preprocess(targets.to(self.device), batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
# 使用 split 方法将预处理后的目标张量 targets 分割为两个部分 : 类别标签 gt_labels 和 边界框坐标 gt_bboxes 。 split((1, 4), 2) 表示沿着第二个维度(索引为2)进行分割,将前1个元素分配给类别标签,剩余的4个元素分配给边界框坐标( xyxy 格式)。
gt_labels, gt_bboxes = targets.split((1, 4), 2) # cls, xyxy
# 计算 边界框坐标 gt_bboxes 在第三个维度(索引为2)上的和,并保留维度。然后使用 gt_(0) 判断和是否大于0,生成一个掩码张量 mask_gt 。这个掩码张量用于指示哪些边界框是有效的(即坐标和大于0)。
mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0)
# 如果在处理目标数据时发生运行时错误( RuntimeError ),则捕获异常并抛出类型错误( TypeError )。错误信息提示数据集格式不正确或不是分割数据集,并建议检查数据集是否为正确格式的分割数据集,例如使用 data=coco8-seg.yaml 。
except RuntimeError as e:
raise TypeError(
"ERROR ❌ segment dataset incorrectly formatted or not a segment dataset.\n" # 错误 ❌ 分割数据集格式不正确或不是分割数据集。
"This error can occur when incorrectly training a 'segment' model on a 'detect' dataset, " # 当在‘检测’数据集上错误地训练‘分割’模型时,
"i.e. 'yolo train model=yolov8n-seg.pt data=coco8.yaml'.\nVerify your dataset is a " # 可能会发生此错误,即‘yolo train model=yolov8n-seg.pt data=coco8.yaml’。
"correctly formatted 'segment' dataset using 'data=coco8-seg.yaml' " # 使用‘data=coco8-seg.yaml’作为示例,验证您的数据集是否为格式正确的‘分割’数据集。
"as an example.\nSee https://docs.ultralytics.com/datasets/segment/ for help." # 请参阅 https://docs.ultralytics.com/datasets/segment/ 获取帮助。
) from e
# 这段代码的主要功能是处理训练批次中的目标数据,包括批量索引、类别标签和边界框坐标。它将这些目标数据进行整合和预处理,生成类别标签张量 gt_labels 、边界框坐标张量 gt_bboxes 和掩码张量 mask_gt 。这些处理后的目标数据将用于后续的损失计算,确保模型能够正确地学习和评估其预测性能。异常处理部分则用于捕获和处理可能发生的错误,确保程序的健壮性和稳定性。
# 这段代码主要负责解码预测的边界框,并使用分配器将预测与目标进行匹配,以便计算损失。
# Pboxes
# 使用 self.bbox_decode 方法将 预测的分布 pred_distri 和 锚点 anchor_points 解码为 预测的边界框 pred_bboxes 。解码后的边界框格式为 xyxy ,形状为 (batch_size, height * width, 4) 。这里 height * width 表示所有锚点的总数,4表示每个边界框的四个坐标(x1, y1, x2, y2)。
pred_bboxes = self.bbox_decode(anchor_points, pred_distri) # xyxy, (b, h*w, 4)
# target_bboxes :匹配后的目标边界框。
# target_scores :匹配后的目标分数。
# fg_mask :前景掩码,表示哪些预测是前景(即目标对象)。
# target_gt_idx :目标索引,表示每个预测对应的目标对象索引。
# 调用 self.assigner 方法将预测的分数、边界框、锚点与目标标签、目标边界框和掩码进行匹配。
# def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt):
# -> 将预测的边界框和分类分数与真实标注进行匹配,并计算目标标签、目标边界框、目标分数等。返回 目标标签 、 目标边界框 、 目标分数 、 前景掩码 和 目标真实标注索引 。
# -> return target_labels, target_bboxes, target_scores, fg_mask.bool(), target_gt_idx
_, target_bboxes, target_scores, fg_mask, target_gt_idx = self.assigner(
# 将预测的分数进行sigmoid激活,转换为概率值,并分离(detach)以避免梯度回传。
pred_scores.detach().sigmoid(),
# 将预测的边界框乘以步长张量,并转换为与目标边界框相同的数据类型,分离以避免梯度回传。
(pred_bboxes.detach() * stride_tensor).type(gt_bboxes.dtype),
# 将锚点乘以步长张量,得到实际的锚点位置。
anchor_points * stride_tensor,
# 目标类别标签。
gt_labels,
# 目标边界框。
gt_bboxes,
# 目标掩码,用于指示有效的目标边界框。
mask_gt,
)
# 计算目标分数的总和,并确保总和至少为1,以避免除以零的情况。这一步是为了在后续的损失计算中使用目标分数的总和进行归一化。
target_scores_sum = max(target_scores.sum(), 1)
# 这段代码解码预测的边界框,将预测的分布和锚点转换为实际的边界框坐标。使用分配器将预测的边界框、分数和锚点与目标数据进行匹配,生成匹配后的目标边界框、目标分数、前景掩码和目标索引。计算目标分数的总和,确保总和至少为1,以便在损失计算中进行归一化。这些步骤为后续的损失计算提供了必要的匹配信息,确保模型能够正确地评估其预测性能。通过这些处理,模型可以更有效地学习和优化其预测结果。
# 这段代码主要负责计算分类损失、边界框损失和掩码损失,并对这些损失进行加权和汇总。最后,返回总损失和各个部分的损失。
# Cls loss
# 注释掉的代码。 self.varifocal_loss 是一个可选的分类损失函数,使用变焦损失(Varifocal Loss)计算分类损失。这里被注释掉了,表示不使用变焦损失。
# loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum # VFL way
# 使用二元交叉熵损失(BCE)计算 分类损失 。 pred_scores 是模型预测的分类分数, target_scores 是目标分数。 target_scores.to(dtype) 确保目标分数的数据类型与预测分数一致。计算损失的总和并除以目标分数的总和进行归一化。
loss[2] = self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum # BCE
# 如果前景掩码 fg_mask 有值,则计算边界框损失。
if fg_mask.sum():
# Bbox loss
# self.bbox_loss 方法接收多个参数,包括预测的分布、预测的边界框、锚点、目标边界框、目标分数、目标分数的总和和前景掩码。返回 边界框损失 loss[0] 和 分布焦点损失 loss[3] 。
loss[0], loss[3] = self.bbox_loss(
pred_distri,
pred_bboxes,
anchor_points,
target_bboxes / stride_tensor,
target_scores,
target_scores_sum,
fg_mask,
)
# Masks loss
# 处理目标掩码 masks ,将其移动到指定设备并转换为浮点类型。
masks = batch["masks"].to(self.device).float()
# 如果掩码的尺寸与原型的尺寸不匹配,则使用最近邻插值进行下采样。
if tuple(masks.shape[-2:]) != (mask_h, mask_w): # downsample
masks = F.interpolate(masks[None], (mask_h, mask_w), mode="nearest")[0]
# 调用 self.calculate_segmentation_loss 方法计算 分割损失 loss[1] 。
loss[1] = self.calculate_segmentation_loss(
fg_mask, masks, target_gt_idx, target_bboxes, batch_idx, proto, pred_masks, imgsz, self.overlap
)
# WARNING: lines below prevent Multi-GPU DDP 'unused gradient' PyTorch errors, do not remove
# 如果前景掩码 fg_mask 没有值,添加虚拟损失以防止多GPU训练中的未使用梯度错误。这一步确保在没有前景目标的情况下,损失计算不会出错。
else:
loss[1] += (proto * 0).sum() + (pred_masks * 0).sum() # inf sums may lead to nan loss
# 对不同部分的损失进行加权。 self.hyp.box 、 self.hyp.cls 和 self.hyp.dfl 是超参数,分别用于调整 边界框损失 、 分割损失 、 分类损失 和 分布焦点损失 的权重。
loss[0] *= self.hyp.box # box gain
loss[1] *= self.hyp.box # seg gain
loss[2] *= self.hyp.cls # cls gain
loss[3] *= self.hyp.dfl # dfl gain
# 计算总损失并乘以批量大小 batch_size ,返回总损失和各个部分的损失。 loss.detach() 确保返回的损失张量不会参与梯度计算,这在某些情况下(如日志记录或评估)是有用的。
return loss.sum() * batch_size, loss.detach() # loss(box, cls, dfl)
# 这段代码计算分类损失,使用二元交叉熵损失(BCE)。如果有前景目标,计算边界框损失和掩码损失。对不同部分的损失进行加权,确保每个部分的损失在总损失中占有适当的权重。返回总损失和各个部分的损失,以便后续的优化和评估。这些步骤确保了模型在训练过程中能够全面地评估其预测性能,并通过多种损失函数进行优化。
# __call__ 方法用于计算分割任务的总损失。它首先初始化损失张量,然后从预测值中提取特征、预测掩码和原型。接着,方法处理目标值,生成锚点,并使用分配器分配目标边界框、分数、前景掩码和目标索引。方法计算分类损失、边界框损失和掩码损失,并对这些损失进行加权。最后,方法返回总损失和各个部分损失的元组。这种方法在处理分割任务时,能够有效地评估模型的预测性能,并通过多种损失函数确保模型的训练效果。
# 这段代码定义了一个静态方法 single_mask_loss ,用于计算单个掩码的损失。
# 装饰器,表示这是一个静态方法,不需要实例化类即可调用。
@staticmethod
# 这行定义了一个名为 single_mask_loss 函数,接受五个参数,返回一个 torch.Tensor 类型的损失值。
# 1.gt_mask :真实掩码,形状为 (n, h, w) ,其中 n 是掩码数量, h 和 w 是掩码的高度和宽度。
# 2.pred :预测掩码的线性组合系数,形状为 (n, m) ,其中 m 是原型张量的通道数。
# 3.proto :原型张量,形状为 (m, h, w) ,其中 m 是通道数, h 和 w 是原型张量的高度和宽度。
# 4.xyxy :目标边界框,形状为 (n, 4) ,每个边界框表示为 (x1, y1, x2, y2) 。
# 5.area :目标边界框的面积,形状为 (n, 1) 。
def single_mask_loss(
gt_mask: torch.Tensor, pred: torch.Tensor, proto: torch.Tensor, xyxy: torch.Tensor, area: torch.Tensor
) -> torch.Tensor:
# 计算单个图像的实例分割损失。
# 注释:
# 该函数使用公式 pred_mask = torch.einsum('in,nhw->ihw', pred, proto) 从原型掩码和预测掩码系数生成预测掩码。
"""
Compute the instance segmentation loss for a single image.
Args:
gt_mask (torch.Tensor): Ground truth mask of shape (n, H, W), where n is the number of objects.
pred (torch.Tensor): Predicted mask coefficients of shape (n, 32).
proto (torch.Tensor): Prototype masks of shape (32, H, W).
xyxy (torch.Tensor): Ground truth bounding boxes in xyxy format, normalized to [0, 1], of shape (n, 4).
area (torch.Tensor): Area of each ground truth bounding box of shape (n,).
Returns:
(torch.Tensor): The calculated mask loss for a single image.
Notes:
The function uses the equation pred_mask = torch.einsum('in,nhw->ihw', pred, proto) to produce the
predicted masks from the prototype masks and predicted mask coefficients.
"""
# 使用爱因斯坦求和约定 torch.einsum 计算预测掩码。这里 pred 的形状为 (n, 32) , proto 的形状为 (32, 80, 80) ,通过矩阵乘法得到预测掩码 pred_mask ,其形状为 (n, 80, 80) 。 n 是预测掩码的数量, 32 是原型的通道数, 80x80 是掩码的空间尺寸。
pred_mask = torch.einsum("in,nhw->ihw", pred, proto) # (n, 32) @ (32, 80, 80) -> (n, 80, 80)
# 使用二元交叉熵损失函数 F.binary_cross_entropy_with_logits 计算预测掩码 pred_mask 和目标掩码 gt_mask 之间的损失。 reduction="none" 表示不对损失进行求和或平均,保留每个元素的损失值。这样可以对损失进行进一步的处理,例如裁剪和归一化。
loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction="none")
# 首先调用 crop_mask 函数对损失张量 loss 进行裁剪,只保留与目标边界框 xyxy 相对应的部分。 crop_mask 函数的作用是将损失张量裁剪到目标掩码的实际区域,避免对背景区域的损失进行计算。
# 然后对裁剪后的损失张量在空间维度(维度1和维度2)上取平均值,得到每个掩码的平均损失。
# 最后将平均损失除以目标掩码的面积 area ,以对损失进行归一化处理。归一化后的损失反映了每个掩码在目标区域内的平均误差。
# 最终对所有掩码的归一化损失求和,得到单个掩码的总损失。
# def crop_mask(masks, boxes): -> 将掩码根据给定的边界框进行裁剪。通过广播运算,将掩码中在边界框范围内的部分保留,范围外的部分置为0,从而实现掩码的裁剪。 -> return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
return (crop_mask(loss, xyxy).mean(dim=(1, 2)) / area).sum()
# 这个静态方法 single_mask_loss 用于计算单个掩码的损失。它通过矩阵乘法生成预测掩码,然后使用二元交叉熵损失函数计算预测掩码和目标掩码之间的损失。为了更准确地反映掩码的实际误差,方法对损失进行裁剪和归一化处理,最终返回单个掩码的总损失。这种方法在处理分割任务时,能够有效地评估模型对目标掩码的预测性能。
# 这段代码定义了 calculate_segmentation_loss 方法,用于计算分割任务的损失。
# 定义了 calculate_segmentation_loss 方法,接收多个参数,用于计算分割任务的损失。
# 1.fg_mask :前景掩码,表示哪些预测是前景(即目标对象)。形状通常为 (batch_size, num_predictions) 。
# 2.masks :目标掩码,表示每个目标对象的实际掩码。形状通常为 (batch_size, height, width) 。
# 3.target_gt_idx :目标索引,表示每个预测对应的目标对象索引。形状通常为 (batch_size, num_predictions) 。
# 4.target_bboxes :目标边界框,表示每个目标对象的边界框坐标(通常为 xyxy 格式)。形状通常为 (batch_size, num_targets, 4) 。
# 5.batch_idx :批量索引,表示每个目标对象所属的批量索引。形状通常为 (num_targets,) 。
# 6.proto :原型张量,用于生成预测掩码。形状通常为 (batch_size, num_prototypes, mask_height, mask_width) 。
# 7.pred_masks :预测掩码,表示模型预测的掩码。形状通常为 (batch_size, num_predictions, num_prototypes) 。
# 8.imgsz :图像尺寸,表示输入图像的尺寸(高度和宽度)。形状通常为 (2,) 。
# 9.overlap :是否重叠的标志,表示目标掩码是否可以重叠。如果为 True ,则目标掩码可以重叠;如果为 False ,则目标掩码不重叠。
def calculate_segmentation_loss(
self,
fg_mask: torch.Tensor,
masks: torch.Tensor,
target_gt_idx: torch.Tensor,
target_bboxes: torch.Tensor,
batch_idx: torch.Tensor,
proto: torch.Tensor,
pred_masks: torch.Tensor,
imgsz: torch.Tensor,
overlap: bool,
) -> torch.Tensor:
# 计算实例分割的损失。
# 注意事项:
# 可以计算批量损失以提高内存使用率的速度。
# 例如,pred_mask 可以按如下方式计算:pred_mask = torch.einsum('in,nhw->ihw', pred, proto) # (i, 32) @ (32, 160, 160) -> (i, 160, 160)
"""
Calculate the loss for instance segmentation.
Args:
fg_mask (torch.Tensor): A binary tensor of shape (BS, N_anchors) indicating which anchors are positive.
masks (torch.Tensor): Ground truth masks of shape (BS, H, W) if `overlap` is False, otherwise (BS, ?, H, W).
target_gt_idx (torch.Tensor): Indexes of ground truth objects for each anchor of shape (BS, N_anchors).
target_bboxes (torch.Tensor): Ground truth bounding boxes for each anchor of shape (BS, N_anchors, 4).
batch_idx (torch.Tensor): Batch indices of shape (N_labels_in_batch, 1).
proto (torch.Tensor): Prototype masks of shape (BS, 32, H, W).
pred_masks (torch.Tensor): Predicted masks for each anchor of shape (BS, N_anchors, 32).
imgsz (torch.Tensor): Size of the input image as a tensor of shape (2), i.e., (H, W).
overlap (bool): Whether the masks in `masks` tensor overlap.
Returns:
(torch.Tensor): The calculated loss for instance segmentation.
Notes:
The batch loss can be computed for improved speed at higher memory usage.
For example, pred_mask can be computed as follows:
pred_mask = torch.einsum('in,nhw->ihw', pred, proto) # (i, 32) @ (32, 160, 160) -> (i, 160, 160)
"""
# 从原型张量 proto 中获取掩码的高度 mask_h 和宽度 mask_w 。
_, _, mask_h, mask_w = proto.shape
# 初始化损失为0,用于累加每个样本的损失。
loss = 0
# Normalize to 0-1
# 将目标边界框 target_bboxes 归一化到0-1范围,使用图像尺寸 imgsz 进行归一化。
target_bboxes_normalized = target_bboxes / imgsz[[1, 0, 1, 0]]
# Areas of target bboxes
# 将归一化后的边界框从 xyxy 格式转换为 xywh 格式,并计算每个边界框的面积。
# def xyxy2xywh(x): -> 将边界框的坐标格式从 xyxy (左上角坐标和右下角坐标)转换为 xywh (中心点坐标、宽度和高度)。返回转换后的坐标张量或数组 y 。 -> return y
marea = xyxy2xywh(target_bboxes_normalized)[..., 2:].prod(2)
# Normalize to mask size
# 将归一化后的边界框坐标转换为掩码尺寸,以便后续处理。
mxyxy = target_bboxes_normalized * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=proto.device)
# 使用 zip 函数将多个张量打包成元组,遍历每个样本的 前景掩码 、 目标索引 、 预测掩码 、 原型 、 归一化边界框 、 面积 和 目标掩码 。
for i, single_i in enumerate(zip(fg_mask, target_gt_idx, pred_masks, proto, mxyxy, marea, masks)):
# 解包当前样本的各个张量。
fg_mask_i, target_gt_idx_i, pred_masks_i, proto_i, mxyxy_i, marea_i, masks_i = single_i
# 检查当前样本的前景掩码是否有值,如果有值则进行损失计算。
if fg_mask_i.any():
# 从目标索引中提取前景掩码对应的索引。
mask_idx = target_gt_idx_i[fg_mask_i]
# 如果 overlap 为 True ,则生成重叠目标掩码;否则,从批量掩码中提取对应的目标掩码。
if overlap:
gt_mask = masks_i == (mask_idx + 1).view(-1, 1, 1)
gt_mask = gt_mask.float()
else:
gt_mask = masks[batch_idx.view(-1) == i][mask_idx]
# 调用 single_mask_loss 方法计算当前样本的分割损失,并累加到总损失中。
loss += self.single_mask_loss(
gt_mask, pred_masks_i[fg_mask_i], proto_i, mxyxy_i[fg_mask_i], marea_i[fg_mask_i]
)
# WARNING: lines below prevents Multi-GPU DDP 'unused gradient' PyTorch errors, do not remove
# 如果前景掩码没有值,添加虚拟损失以防止多GPU训练中的未使用梯度错误。
else:
loss += (proto * 0).sum() + (pred_masks * 0).sum() # inf sums may lead to nan loss
# 返回总损失除以前景掩码的总和,得到平均分割损失。
return loss / fg_mask.sum()
# calculate_segmentation_loss 方法用于计算分割任务的损失。它首先对目标边界框进行归一化和尺寸转换,然后遍历每个样本,生成目标掩码,并调用 single_mask_loss 方法计算每个样本的分割损失。最后,方法返回所有样本的平均分割损失。这种方法在处理分割任务时,能够有效地评估模型对目标掩码的预测性能,并通过归一化和裁剪处理,确保损失计算的准确性和稳定性。
# 这个类 v8SegmentationLoss 用于计算分割任务的训练损失,包括框损失、分类损失、分布损失和掩码损失。通过继承 v8DetectionLoss 类,它复用了检测任务的一些功能,并增加了掩码损失的计算。
# 这段代码定义了一个名为 v8PoseLoss 的类,它是 v8DetectionLoss 类的子类,用于计算训练过程中的损失函数,特别是在目标检测和人体姿态估计任务中。
class v8PoseLoss(v8DetectionLoss):
# 计算训练损失的标准类。
"""Criterion class for computing training losses."""
# 这段代码是 v8PoseLoss 类的初始化方法 __init__ ,用于设置损失计算所需的参数和对象。
# 定义 v8PoseLoss 类的初始化方法 __init__ ,接受一个参数。
# 1.model :这个参数是一个模型对象。注释中提到 model 必须是去并行化的,这意味着在多GPU训练时,模型应该已经被包装在 DataParallel 或 DistributedDataParallel 之外。
def __init__(self, model): # model must be de-paralleled
# 使用模型初始化 v8PoseLoss,设置关键点变量并声明关键点损失实例。
"""Initializes v8PoseLoss with model, sets keypoint variables and declares a keypoint loss instance."""
# 调用父类 v8DetectionLoss 的初始化方法,将 model 传递给父类的构造函数。这一步确保了父类的初始化逻辑得以执行,例如设置一些通用的属性或初始化通用的损失函数。
super().__init__(model)
# 从模型的最后一个模块中获取 关键点的形状 ( kpt_shape ),并将其存储在 self.kpt_shape 中。这个形状通常是一个列表,例如 [17, 3] ,表示有17个关键点,每个关键点有3个坐标值(通常是x, y, 可见性)。
self.kpt_shape = model.model[-1].kpt_shape
# 初始化一个二元交叉熵损失函数( BCEWithLogitsLoss ),用于计算关键点的可见性损失。这个损失函数适用于二分类问题,并且在内部应用了sigmoid激活函数,因此输入应该是未经sigmoid处理的logits。
self.bce_pose = nn.BCEWithLogitsLoss()
# 检查关键点的形状是否为 [17, 3] ,并将结果存储在 is_pose 变量中。这个变量用于后续判断是否使用预定义的关键点标准差( OKS_SIGMA )。
is_pose = self.kpt_shape == [17, 3]
# 从 kpt_shape 中提取关键点的数量,并将其存储在 nkpt 变量中。例如,如果 kpt_shape 是 [17, 3] ,那么 nkpt 将是17。
nkpt = self.kpt_shape[0] # number of keypoints
# 如果 is_pose 为 True ,则从 OKS_SIGMA (一个预定义的numpy数组,包含关键点的标准差,用于计算OKS,即目标关键点相似度)中创建一个张量,并将其移动到当前设备( self.device )。
# 如果 is_pose 为 False ,则创建一个全1的张量,其形状为 (nkpt,) ,并除以 nkpt ,以得到均匀的权重。
# 这个张量 sigmas 用于后续的关键点损失计算中,为不同的关键点分配不同的权重。
sigmas = torch.from_numpy(OKS_SIGMA).to(self.device) if is_pose else torch.ones(nkpt, device=self.device) / nkpt
# 使用 sigmas 初始化一个 KeypointLoss 实例,并将其存储在 self.keypoint_loss 中。 KeypointLoss 是一个自定义的损失函数类,专门用于计算关键点的位置损失。通过传递 sigmas ,可以为不同的关键点设置不同的权重,这对于人体姿态估计等任务中不同关键点的重要性不同是有帮助的。
self.keypoint_loss = KeypointLoss(sigmas=sigmas)
# 这段初始化代码的主要目的是设置 v8PoseLoss 类所需的各种属性和损失函数。关键步骤包括 :继承父类并调用其初始化方法。从模型中获取关键点的形状,并据此判断是否使用预定义的关键点标准差。初始化用于计算关键点可见性损失的二元交叉熵损失函数。根据关键点的形状和是否使用预定义标准差,初始化关键点位置损失函数。通过这些步骤, v8PoseLoss 类能够为后续的损失计算准备好必要的工具和参数,特别是在处理关键点检测和姿态估计任务时。
# 这段代码定义了 v8PoseLoss 类的 __call__ 方法,用于计算模型的总损失。这个方法将预测值和批次数据作为输入,计算并返回总损失以及各个分项损失。
# __call__ 方法定义了如何调用 v8PoseLoss 实例。
# 1.preds :模型的预测输出。
# 2.batch :当前批次的数据。
def __call__(self, preds, batch):
# 计算总损失并将其分离。
"""Calculate the total loss and detach it."""
# 这段代码是 v8PoseLoss 类的 __call__ 方法的一部分,用于初始化损失张量并处理模型的预测输出。
# 初始化一个零张量 loss ,形状为 (5,) ,用于存储五个分项损失。 box :框损失。 cls :分类损失。 dfl :分布焦点损失。 kpt_location :关键点位置损失。 kpt_visibility :关键点可见性损失。 device=self.device 确保 loss 张量在正确的设备上(例如CPU或GPU)。
loss = torch.zeros(5, device=self.device) # box, cls, dfl, kpt_location, kpt_visibility
# 从 preds 中提取特征图 feats 和 预测关键点 pred_kpts 。 preds 可能是一个列表或一个元组。如果 preds[0] 是一个列表,直接使用 preds ;否则,使用 preds[1] 。这种设计是为了处理不同格式的预测输出。
feats, pred_kpts = preds if isinstance(preds[0], list) else preds[1]
# feats 是一个包含多个特征图的列表。使用列表推导式将每个特征图 xi 重新排列为 (batch_size, num_outputs, -1) 的形状,然后使用 torch.cat 沿第2维(特征维)拼接这些特征图。
# split((self.reg_max * 4, self.nc), 1) 将拼接后的张量分割为两个部分 : pred_distri 和 pred_scores 。
# pred_distri 的形状为 (batch_size, num_anchors, self.reg_max * 4) ,表示 预测的分布 ,用于解码为框的坐标。
# pred_scores 的形状为 (batch_size, num_anchors, self.nc) ,表示 预测的分类分数 。
# self.reg_max 是一个超参数,表示分布的最大值。 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
)
# 这段代码的主要目的是初始化损失张量,并从模型的预测输出中提取和处理特征图和预测关键点。初始化一个零张量 loss ,用于存储五个分项损失。从 preds 中提取特征图 feats 和预测关键点 pred_kpts 。将特征图重新排列并拼接,然后分割为预测分布 pred_distri 和预测分数 pred_scores 。通过这些步骤,可以为后续的损失计算准备必要的数据。这些数据将用于评估模型在目标检测和关键点检测任务中的性能。
# 这段代码对预测的分数、分布和关键点进行维度重排,以确保它们的形状适合后续的处理。同时,它还计算了图像的大小和生成了锚点。
# B, grids, ..
# permute(0, 2, 1) 重新排列 pred_scores 的维度,使其形状从 (batch_size, num_anchors, num_classes) 变为 (batch_size, num_classes, num_anchors) 。这种重排是为了后续计算分类损失时更方便地处理类别维度。
# contiguous() 返回一个新的张量,该张量在内存中是连续的。这通常在进行维度重排后调用,以确保张量在后续操作中可以高效地使用。
pred_scores = pred_scores.permute(0, 2, 1).contiguous()
# 对 pred_distri 进行维度重排和连续化操作,使其形状从 (batch_size, num_anchors, self.reg_max * 4) 变为 (batch_size, self.reg_max * 4, num_anchors) 。这种重排是为了后续解码框坐标时更方便地处理分布维度。
pred_distri = pred_distri.permute(0, 2, 1).contiguous()
# 对 pred_kpts 进行的维度重排和连续化操作,使其形状从 (batch_size, num_anchors, num_keypoints, 3) 变为 (batch_size, num_keypoints, 3, num_anchors) 。这种重排是为了后续处理关键点时更方便地处理关键点维度。
pred_kpts = pred_kpts.permute(0, 2, 1).contiguous()
# 获取 pred_scores 的数据类型 dtype ,以便在后续操作中使用相同的数据类型,确保数据的一致性。
dtype = pred_scores.dtype
# 计算图像的大小 imgsz 。 feats[0].shape[2:] 获取第一个特征图的空间维度(高度和宽度),然后将其转换为张量并乘以步长 self.stride[0] ,得到图像的实际大小。
# device=self.device 确保 imgsz 在正确的设备上(例如CPU或GPU)。 dtype=dtype 确保 imgsz 的数据类型与 pred_scores 相同。
imgsz = torch.tensor(feats[0].shape[2:], device=self.device, dtype=dtype) * self.stride[0] # image size (h,w)
# 调用 make_anchors 函数生成锚点 anchor_points 和步长张量 stride_tensor 。 feats 是特征图的列表, self.stride 是步长列表, 0.5 是一个偏移量参数。这个函数根据特征图和步长生成锚点,用于后续的框解码和损失计算。
# def make_anchors(feats, strides, grid_cell_offset=0.5): -> 用于生成目标检测模型中使用的锚点(anchors)。返回连接后的 锚点坐标张量 和 步长张量 。 -> return torch.cat(anchor_points), torch.cat(stride_tensor)
anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5)
# 这段代码的主要目的是对预测的分数、分布和关键点进行维度重排,以确保它们的形状适合后续的处理。同时,它还计算了图像的大小和生成了锚点。对 pred_scores 、 pred_distri 和 pred_kpts 进行维度重排和连续化操作,使其形状适合后续的处理。获取 pred_scores 的数据类型 dtype 。计算图像的大小 imgsz 。调用 make_anchors 函数生成锚点 anchor_points 和步长张量 stride_tensor 。通过这些步骤,可以为后续的损失计算和框解码等操作准备必要的数据结构和参数。
# 这段代码的目的是处理和准备目标(ground truth)数据,以便用于后续的损失计算。具体步骤包括提取批次大小、创建批次索引、拼接目标数据、预处理目标数据、分离类别标签和框坐标,以及生成掩码。
# Targets
# 从 pred_scores 中提取批次大小 batch_size 。 pred_scores 的形状为 (batch_size, num_classes, num_anchors) ,因此 pred_scores.shape[0] 给出了批次大小。
batch_size = pred_scores.shape[0]
# 从 batch 字典中提取 batch_idx ,这是一个表示每个目标属于哪个图像的索引。使用 view(-1, 1) 将其形状调整为 (N, 1) ,其中 N 是目标的总数。
batch_idx = batch["batch_idx"].view(-1, 1)
# 使用 torch.cat 沿第1维(列)拼接 batch_idx 、 类别标签 batch["cls"] 和 真实框 batch["bboxes"] 。 batch["cls"].view(-1, 1) 将类别标签的形状调整为 (N, 1) , batch["bboxes"] 的形状为 (N, 4) ,因此 targets 的形状为 (N, 6) ,其中每一行表示一个目标的 批次索引 、 类别标签 和 框坐标 。
targets = torch.cat((batch_idx, batch["cls"].view(-1, 1), batch["bboxes"]), 1)
# 将 targets 移动到当前设备(例如CPU或GPU)。 调用 self.preprocess 方法对 targets 进行预处理。 imgsz[[1, 0, 1, 0]] 用于调整图像大小的顺序,确保框坐标的顺序为 (x1, y1, x2, y2) 。
targets = self.preprocess(targets.to(self.device), batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
# 使用 split 方法将 targets 分离为 类别标签 gt_labels 和 真实框 gt_bboxes 。 split((1, 4), 2) 表示在第2维(列)上将 targets 分为两部分,第一部分包含1列(类别标签),第二部分包含4列(框坐标)。
gt_labels, gt_bboxes = targets.split((1, 4), 2) # cls, xyxy
# 计算 gt_bboxes 的每一行的和,然后生成一个掩码 mask_gt ,表示每个框是否有效。 sum(2, keepdim=True) 沿第2维(列)计算和,保持输出的维度为 (N, 1) 。 gt_(0) 生成一个布尔掩码,表示每个框的和是否大于0,即框是否有效。
# sum(2, keepdim=True) 沿第2维(列)计算每个框的坐标和。 keepdim=True 保持输出的维度为 (N, 1) ,这样每个框的坐标和是一个标量值。
# 例如,如果 gt_bboxes 为 :
# tensor([[10, 10, 20, 20],
# [0, 0, 0, 0],
# [15, 15, 30, 30]])
# 那么 gt_bboxes.sum(2, keepdim=True) 的结果为 :
# tensor([[60],
# [0],
# [90]])
mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0)
# 这段代码的主要目的是处理和准备目标数据,以便用于后续的损失计算。提取批次大小 batch_size 。创建批次索引 batch_idx 。拼接目标数据,包括批次索引、类别标签和真实框。对目标数据进行预处理,调整框坐标的尺度。分离类别标签和真实框。生成掩码 mask_gt ,表示每个框是否有效。通过这些步骤,可以确保目标数据的格式和内容适合用于后续的损失计算,从而评估模型在目标检测任务中的性能。
# 这段代码的目的是解码预测的框和关键点,并使用分配器(assigner)将预测与真实目标进行匹配。
# Pboxes
# 调用 bbox_decode 方法,将预测的分布 pred_distri 和锚点 anchor_points 解码为预测的框坐标 pred_bboxes 。 pred_bboxes 的形状为 (batch_size, h*w, 4) ,其中 h*w 是特征图上的锚点数量,4 表示每个框的坐标 (x1, y1, x2, y2) 。
pred_bboxes = self.bbox_decode(anchor_points, pred_distri) # xyxy, (b, h*w, 4)
# pred_kpts.view(batch_size, -1, *self.kpt_shape) 将预测的关键点 pred_kpts 重新排列为 (batch_size, h*w, num_keypoints, 3) 的形状。
# self.kpts_decode(anchor_points, ...) 调用 kpts_decode 方法,将预测的关键点坐标从相对锚点的位置解码为绝对位置。
# pred_kpts 的形状为 (batch_size, h*w, 17, 3) ,其中 17 是关键点的数量,3 表示每个关键点的坐标 (x, y, visibility) 。
pred_kpts = self.kpts_decode(anchor_points, pred_kpts.view(batch_size, -1, *self.kpt_shape)) # (b, h*w, 17, 3)
# self.assigner(...) 调用分配器方法,将预测的框和关键点与真实目标进行匹配。分配器返回以下值 : _ 忽略的返回值。 target_bboxes 与预测匹配的真实框坐标。 target_scores 与预测匹配的真实目标的分数。 fg_mask 前景掩码,表示哪些预测是前景(即与真实目标匹配)。 target_gt_idx 与预测匹配的真实目标的索引。
# def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt):
# -> 将预测的边界框和分类分数与真实标注进行匹配,并计算目标标签、目标边界框、目标分数等。返回 目标标签 、 目标边界框 、 目标分数 、 前景掩码 和 目标真实标注索引 。
# -> return target_labels, target_bboxes, target_scores, fg_mask.bool(), target_gt_idx
_, target_bboxes, target_scores, fg_mask, target_gt_idx = self.assigner(
# 计算预测分数的 sigmoid 值,将其从 logits 转换为概率。
pred_scores.detach().sigmoid(),
# 将预测的框坐标 pred_bboxes 乘以步长 stride_tensor ,并转换为与 gt_bboxes 相同的数据类型。
(pred_bboxes.detach() * stride_tensor).type(gt_bboxes.dtype),
# 计算锚点的绝对位置。
anchor_points * stride_tensor,
gt_labels,
gt_bboxes,
mask_gt,
)
# 这段代码的主要目的是解码预测的框和关键点,并使用分配器将预测与真实目标进行匹配。解码预测的框坐标 pred_bboxes 。解码预测的关键点坐标 pred_kpts 。使用分配器将预测的框和关键点与真实目标进行匹配,生成匹配结果和相关掩码。通过这些步骤,可以为后续的损失计算准备必要的匹配信息,确保模型的训练能够正确地评估预测的准确性和有效性。
# 这段代码的目的是计算分类损失和框损失,并在有前景掩码的情况下计算关键点损失。
# 计算 target_scores 的总和,确保其不为零。 max(target_scores.sum(), 1) 用于避免除以零的情况,确保损失计算的稳定性。
target_scores_sum = max(target_scores.sum(), 1)
# Cls loss
# loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum # VFL way
# 计算分类损失。这里使用了二元交叉熵损失(BCE)而不是变焦损失(Varifocal Loss, VFL)。
# pred_scores 是预测的分类分数, target_scores 是目标分数。 target_scores.to(dtype) 确保目标分数的数据类型与预测分数相同。
# self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum 计算二元交叉熵损失的总和,并除以 target_scores_sum 进行归一化。
loss[3] = self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum # BCE
# Bbox loss
# 如果前景掩码 fg_mask 的总和大于零,表示有有效的前景预测。
if fg_mask.sum():
# 将目标框坐标除以步长张量,转换为相对坐标。
target_bboxes /= stride_tensor
# 计算框损失,返回 框损失 loss[0] 和 分布焦点损失 loss[4] 。
loss[0], loss[4] = self.bbox_loss(
pred_distri, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask
)
# 从 batch 中提取关键点数据 keypoints ,并将其移动到当前设备,转换为浮点数,并克隆。
keypoints = batch["keypoints"].to(self.device).float().clone()
# 将关键点的 x 和 y 坐标乘以图像的宽度和高度,转换为绝对坐标。
keypoints[..., 0] *= imgsz[1]
keypoints[..., 1] *= imgsz[0]
# 调用 self.calculate_keypoints_loss(...) 计算 关键点位置损失 loss[1] 和 关键点对象损失 loss[2] 。
loss[1], loss[2] = self.calculate_keypoints_loss(
# fg_mask 是前景掩码, target_gt_idx 是目标索引, keypoints 是真实关键点, batch_idx 是批次索引, stride_tensor 是步长张量, target_bboxes 是目标框, pred_kpts 是预测关键点。
fg_mask, target_gt_idx, keypoints, batch_idx, stride_tensor, target_bboxes, pred_kpts
)
# 这段代码的主要目的是计算分类损失、框损失和关键点损失。计算目标分数的总和,确保其不为零。计算分类损失,使用二元交叉熵损失。如果有前景掩码,计算框损失和分布焦点损失。将真实关键点的坐标转换为绝对坐标。计算关键点位置损失和关键点对象损失。通过这些步骤,可以全面评估模型在目标检测和关键点检测任务中的性能,并为反向传播提供必要的梯度信息。
# 这段代码的目的是将各个分项损失乘以相应的超参数权重,然后计算总损失并返回。这些超参数权重用于平衡不同损失项在总损失中的贡献。
# 将框损失 loss[0] 乘以超参数 self.hyp.box ,这个超参数用于调整 框损失 在总损失中的权重。
loss[0] *= self.hyp.box # box gain
# 将关键点位置损失 loss[1] 乘以超参数 self.hyp.pose ,这个超参数用于调整 关键点位置损失 在总损失中的权重。
loss[1] *= self.hyp.pose # pose gain
# 将关键点对象损失 loss[2] 乘以超参数 self.hyp.kobj ,这个超参数用于调整 关键点对象损失 在总损失中的权重。
loss[2] *= self.hyp.kobj # kobj gain
# 将分类损失 loss[3] 乘以超参数 self.hyp.cls ,这个超参数用于调整 分类损失 在总损失中的权重。
loss[3] *= self.hyp.cls # cls gain
# 将分布焦点损失 loss[4] 乘以超参数 self.hyp.dfl ,这个超参数用于调整 分布焦点损失 在总损失中的权重。
loss[4] *= self.hyp.dfl # dfl gain
# 计算总损失 loss.sum() ,并将其乘以批次大小 batch_size ,以得到整个批次的总损失。 使用 loss.detach() 返回各个分项损失的副本,这些副本不会参与梯度计算,通常用于日志记录或监控。
# 返回值是一个元组,包含总损失和各个分项损失。
return loss.sum() * batch_size, loss.detach() # loss(box, cls, dfl)
# 这段代码的主要目的是将各个分项损失乘以相应的超参数权重,然后计算总损失并返回。将框损失乘以 self.hyp.box 。将关键点位置损失乘以 self.hyp.pose 。将关键点对象损失乘以 self.hyp.kobj 。将分类损失乘以 self.hyp.cls 。将分布焦点损失乘以 self.hyp.dfl 。计算总损失,并将其乘以批次大小。返回总损失和各个分项损失的副本。通过这些步骤,可以确保不同损失项在总损失中的贡献得到合理调整,从而更好地平衡模型在不同任务上的性能。这对于模型的训练和优化非常重要。
# 这段代码的主要目的是计算模型的总损失,包括框损失、分类损失、分布焦点损失、关键点位置损失和关键点对象损失。初始化损失张量。提取和处理预测值和批次数据。生成锚点和步长张量。预处理目标数据。解码预测框和预测关键点。使用 assigner 方法匹配预测值和真实值。计算分类损失。如果有前景掩码,计算框损失和分布焦点损失。计算关键点位置损失和关键点对象损失。将各个分项损失乘以相应的权重。返回总损失和各个分项损失。通过这些步骤,可以全面评估模型在目标检测和关键点检测任务中的性能,并为反向传播提供必要的梯度信息。
# 这段代码定义了一个静态方法 kpts_decode ,用于将预测的关键点坐标从相对坐标转换为绝对坐标。
# 使用 @staticmethod 装饰器定义一个静态方法 kpts_decode 。静态方法不需要类实例即可调用,通常用于与类相关的功能,但不依赖于类的实例属性。
@staticmethod
# 这行定义了一个名为 kpts_decode 的方法,接受两个参数。
# 1.anchor_points :锚点坐标。
# 2.pred_kpts :预测的关键点坐标。
def kpts_decode(anchor_points, pred_kpts):
# 将预测的关键点解码为图像坐标。
"""Decodes predicted keypoints to image coordinates."""
# 创建 pred_kpts 的副本并存储在变量 y 中。这一步是为了避免修改原始的 pred_kpts 数据。
y = pred_kpts.clone()
# 将 y 中每个关键点的前两个坐标(x和y)乘以2.0。这一步通常是为了将相对坐标转换为更大的范围,以便后续的调整。
y[..., :2] *= 2.0
# 将 y 中每个关键点的x坐标加上锚点的x坐标减去0.5。这一步是为了将关键点的x坐标从相对锚点的位置转换为绝对位置。
y[..., 0] += anchor_points[:, [0]] - 0.5
# 将 y 中每个关键点的y坐标加上锚点的y坐标减去0.5。这一步是为了将关键点的y坐标从相对锚点的位置转换为绝对位置。
y[..., 1] += anchor_points[:, [1]] - 0.5
# 返回转换后的关键点坐标 y 。
return y
# kpts_decode 方法的主要目的是将预测的关键点坐标从相对锚点的位置转换为绝对位置。创建预测关键点坐标的副本。将每个关键点的x和y坐标乘以2.0,扩大坐标范围。将每个关键点的x坐标加上锚点的x坐标减去0.5,调整x坐标。将每个关键点的y坐标加上锚点的y坐标减去0.5,调整y坐标。通过这些步骤, kpts_decode 方法能够将预测的关键点坐标转换为图像中的绝对位置,这对于后续的损失计算和关键点检测任务非常重要。
# 这段代码定义了一个方法 calculate_keypoints_loss ,用于计算模型的关键点损失和关键点对象损失。这个方法的目的是评估预测的关键点与真实关键点之间的差异,并计算关键点的可见性损失。
# 定义 calculate_keypoints_loss 方法,接受以下参数 :
# 1.masks :二元掩码张量,表示对象的存在性,形状为 (BS, N_anchors) 。
# 2.target_gt_idx :索引张量,将锚点映射到真实对象,形状为 (BS, N_anchors) 。
# 3.keypoints :真实关键点,形状为 (N_kpts_in_batch, N_kpts_per_object, kpts_dim) 。
# 4.batch_idx :关键点的批次索引,形状为 (N_kpts_in_batch, 1) 。
# 5.stride_tensor :锚点的步长张量,形状为 (N_anchors, 1) 。
# 6.target_bboxes :真实框,格式为 (x1, y1, x2, y2) ,形状为 (BS, N_anchors, 4) 。
# 7.pred_kpts :预测的关键点,形状为 (BS, N_anchors, N_kpts_per_object, kpts_dim) 。
def calculate_keypoints_loss(
self, masks, target_gt_idx, keypoints, batch_idx, stride_tensor, target_bboxes, pred_kpts
):
# 计算模型的关键点损失。
# 此函数计算给定批次的关键点损失和关键点对象损失。关键点损失基于预测的关键点和地面实况关键点之间的差异。关键点对象损失是二元分类损失,用于对关键点是否存在进行分类。
"""
Calculate the keypoints loss for the model.
This function calculates the keypoints loss and keypoints object loss for a given batch. The keypoints loss is
based on the difference between the predicted keypoints and ground truth keypoints. The keypoints object loss is
a binary classification loss that classifies whether a keypoint is present or not.
Args:
masks (torch.Tensor): Binary mask tensor indicating object presence, shape (BS, N_anchors).
target_gt_idx (torch.Tensor): Index tensor mapping anchors to ground truth objects, shape (BS, N_anchors).
keypoints (torch.Tensor): Ground truth keypoints, shape (N_kpts_in_batch, N_kpts_per_object, kpts_dim).
batch_idx (torch.Tensor): Batch index tensor for keypoints, shape (N_kpts_in_batch, 1).
stride_tensor (torch.Tensor): Stride tensor for anchors, shape (N_anchors, 1).
target_bboxes (torch.Tensor): Ground truth boxes in (x1, y1, x2, y2) format, shape (BS, N_anchors, 4).
pred_kpts (torch.Tensor): Predicted keypoints, shape (BS, N_anchors, N_kpts_per_object, kpts_dim).
Returns:
(tuple): Returns a tuple containing:
- kpts_loss (torch.Tensor): The keypoints loss.
- kpts_obj_loss (torch.Tensor): The keypoints object loss.
"""
# 这段代码的作用是将关键点数据按照批次组织起来,以便后续能够根据每个图像中的关键点进行进一步的处理。
# 将 batch_idx 展平成一维张量。 batch_idx 通常是一个二维张量,其中包含每个关键点所属的批次索引。展平后,它变成一个一维张量,每个元素表示一个关键点所属的批次索引。
batch_idx = batch_idx.flatten()
# 获取批次大小 batch_size 。这里假设 masks 的长度与批次大小相同,因为 masks 通常是一个二维张量,其中每一行对应一个图像的掩码。
batch_size = len(masks)
# Find the maximum number of keypoints in a single image
# 计算单个图像中最大关键点数量 max_kpts 。 torch.unique(batch_idx, return_counts=True) 返回一个元组,其中第一个元素是 batch_idx 中的唯一值(即不同批次的索引),第二个元素是每个唯一值出现的次数(即每个批次中的关键点数量)。 [1].max() 获取这些次数中的最大值,即单个图像中最大关键点数量。
max_kpts = torch.unique(batch_idx, return_counts=True)[1].max()
# Create a tensor to hold batched keypoints
# 创建一个零张量 batched_keypoints ,用于存储每个图像的关键点。
# 形状为 (batch_size, max_kpts, keypoints.shape[1], keypoints.shape[2]) ,其中 batch_size 是批次大小, max_kpts 是单个图像中最大关键点数量, keypoints.shape[1] 和 keypoints.shape[2] 分别是 关键点的维度 和 特征维度 。 device=keypoints.device 确保 batched_keypoints 与 keypoints 在同一设备上(例如CPU或GPU)。
batched_keypoints = torch.zeros(
(batch_size, max_kpts, keypoints.shape[1], keypoints.shape[2]), device=keypoints.device
)
# TODO: any idea how to vectorize this? TODO:知道如何将其矢量化吗?
# Fill batched_keypoints with keypoints based on batch_idx
# 根据 batch_idx 填充 batched_keypoints 。
# 遍历每个批次索引。
for i in range(batch_size):
# 从 keypoints 中选择属于当前批次索引 i 的关键点。
keypoints_i = keypoints[batch_idx == i]
# 将选择的关键点 keypoints_i 填充到 batched_keypoints 中对应的位置。 这里使用了循环来实现,注释中提到“TODO: any idea how to vectorize this?”,意味着作者在寻找一种向量化的方法来替代这个循环,以提高效率。
# 这行代码是将属于第 i 个图像的关键点数据 keypoints_i 填充到 batched_keypoints 张量中对应的位置。
# batched_keypoints 是一个零张量,其形状为 (batch_size, max_kpts, keypoints.shape[1], keypoints.shape[2]) 。这个张量用于存储每个图像的关键点数据,其中 batch_size 是批次大小, max_kpts 是单个图像中最大关键点数量, keypoints.shape[1] 和 keypoints.shape[2] 分别是 关键点的维度 和 特征维度 。
# i 是当前处理的图像索引。
# keypoints_i 是属于第 i 个图像的关键点数据,其形状为 (num_kpts_i, keypoints.shape[1], keypoints.shape[2]) ,其中 num_kpts_i 是第 i 个图像中的关键点数量。 keypoints_i.shape[0] 是第 i 个图像中的关键点数量 num_kpts_i 。
# batched_keypoints[i, : keypoints_i.shape[0]] 表示 batched_keypoints 中第 i 个图像的前 num_kpts_i 个关键点的位置。
# = 操作符用于将 keypoints_i 的内容复制到 batched_keypoints 的相应位置。
# 选择目标位置 : batched_keypoints[i, : keypoints_i.shape[0]] 选择了 batched_keypoints 中第 i 个图像的前 num_kpts_i 个关键点的位置。这里使用了切片操作 : 来选择前 num_kpts_i 个关键点。
# 赋值操作 : keypoints_i 包含了第 i 个图像的所有关键点数据。 = 操作符将 keypoints_i 的内容复制到 batched_keypoints 的相应位置,从而将第 i 个图像的关键点数据填充到 batched_keypoints 中。
# 这行代码的作用是将每个图像的关键点数据填充到 batched_keypoints 张量中对应的位置。通过这种方式,可以将所有图像的关键点数据组织成一个统一的张量,便于后续的处理和计算。
batched_keypoints[i, : keypoints_i.shape[0]] = keypoints_i
# 这段代码的主要目的是将关键点数据按照批次组织起来,以便后续能够根据每个图像中的关键点进行进一步的处理。将批次索引 batch_idx 展平。获取批次大小 batch_size 。计算单个图像中最大关键点数量 max_kpts 。创建一个零张量 batched_keypoints ,用于存储每个图像的关键点。根据批次索引 batch_idx 填充 batched_keypoints 。通过这些步骤,可以将关键点数据按批次组织,为后续的关键点损失计算和其他处理提供了便利。
# 这段代码的目的是从 batched_keypoints 中选择 与预测锚点匹配的真实关键点 ,并将这些关键点的坐标转换为绝对坐标。最后,初始化关键点位置损失和关键点对象损失。
# Expand dimensions of target_gt_idx to match the shape of batched_keypoints
# target_gt_idx 是一个形状为 (BS, N_anchors) 的张量,表示每个锚点对应的真实对象索引。
# unsqueeze(-1).unsqueeze(-1) 用于扩展 target_gt_idx 的维度,使其形状变为 (BS, N_anchors, 1, 1) 。这样可以使其维度与 batched_keypoints 的维度 (batch_size, max_kpts, keypoints.shape[1], keypoints.shape[2]) 兼容。
target_gt_idx_expanded = target_gt_idx.unsqueeze(-1).unsqueeze(-1)
# Use target_gt_idx_expanded to select keypoints from batched_keypoints
# torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor
# torch.gather 是 PyTorch 中的一个函数,它用于根据索引从输入张量中抽取值。这个函数在处理数据时提供了极大的灵活性,特别是在深度学习任务中,它允许我们动态地选择和重新排序张量中的元素。
# 参数 :
# input (Tensor) :目标变量,即要从中收集值的输入张量。
# dim (int) :需要沿着取值的坐标轴(维度)。这意味着你希望沿着哪个维度根据索引收集值。
# index (LongTensor) :包含要收集的元素索引的张量。它的 shape 应该与 input 在非 dim 维度上的形状相同,而在 dim 维度上,其数值表示 input 中需要被收集的数据的索引位置。
# sparse_grad (bool, optional) :如果为真,输入将是一个稀疏张量。默认值为 False 。
# out (Tensor, optional) :输出张量。如果提供,输出结果将被写入此张量。默认值为 None 。
# 返回值 :
# Tensor :一个新的张量,包含了根据 index 从 input 张量中收集的值。
# 工作原理 :
# torch.gather 函数的核心作用是从输入张量 input 中按照指定的索引 index 抽取值。在指定的维度 dim 上, index 张量中的每个值指示了要从 input 中抽取哪个位置的值。
# torch.gather 函数在深度学习中非常有用,特别是在需要根据目标标签从输出的概率分布中直接获取相应的概率值时,这在计算损失函数时尤其有用。
# target_gt_idx_expanded.expand(-1, -1, keypoints.shape[1], keypoints.shape[2]) 将 target_gt_idx_expanded 的形状扩展为 (BS, N_anchors, keypoints.shape[1], keypoints.shape[2]) ,以匹配 batched_keypoints 的形状。
# batched_keypoints.gather(1, ...) 使用扩展后的 target_gt_idx_expanded 从 batched_keypoints 中选择与预测锚点匹配的真实关键点。 gather 方法的第一个参数 1 表示在第1维(即 max_kpts 维)上进行选择。
selected_keypoints = batched_keypoints.gather(
1, target_gt_idx_expanded.expand(-1, -1, keypoints.shape[1], keypoints.shape[2])
)
# Divide coordinates by stride
# stride_tensor 是一个形状为 (N_anchors, 1) 的张量,表示每个锚点的步长。
# stride_tensor.view(1, -1, 1, 1) 用于调整 stride_tensor 的形状,使其变为 (1, N_anchors, 1, 1) ,以匹配 selected_keypoints 的形状。
# selected_keypoints /= stride_tensor.view(1, -1, 1, 1) 将 selected_keypoints 中的坐标除以步长,将相对坐标转换为绝对坐标。
selected_keypoints /= stride_tensor.view(1, -1, 1, 1)
# 初始化 关键点位置损失 kpts_loss 和 关键点对象损失 kpts_obj_loss 为0。这些变量将用于后续计算损失值。
kpts_loss = 0
kpts_obj_loss = 0
# 这段代码的主要目的是从 batched_keypoints 中选择与预测锚点匹配的真实关键点,并将这些关键点的坐标转换为绝对坐标。扩展 target_gt_idx 的维度,使其形状与 batched_keypoints 兼容。使用扩展后的 target_gt_idx 从 batched_keypoints 中选择与预测锚点匹配的真实关键点。将选择的关键点坐标除以步长,转换为绝对坐标。初始化关键点位置损失和关键点对象损失。通过这些步骤,可以为后续的关键点损失计算准备数据。
# 这段代码用于计算关键点位置损失和关键点对象损失。
# 检查 masks 张量中是否有任何 True 值。如果 masks 中至少有一个 True 值,表示有有效的预测锚点,那么执行以下步骤。
if masks.any():
# 从 selected_keypoints 中选择与 masks 中 True 值对应的关键点。 gt_kpt 是这些真实关键点的集合。
gt_kpt = selected_keypoints[masks]
# target_bboxes[masks] 从 target_bboxes 中选择与 masks 中 True 值对应的框。
# xyxy2xywh(...) 将框的格式从 (x1, y1, x2, y2) 转换为 (x_center, y_center, width, height) 。
# [:, 2:].prod(1, keepdim=True) 计算每个框的面积,即 width * height ,并保持输出的维度为 (N, 1) ,其中 N 是有效框的数量。
# def xyxy2xywh(x): -> 将边界框的坐标格式从 xyxy (左上角坐标和右下角坐标)转换为 xywh (中心点坐标、宽度和高度)。返回转换后的坐标张量或数组 y 。 -> return y
area = xyxy2xywh(target_bboxes[masks])[:, 2:].prod(1, keepdim=True)
# 从 pred_kpts 中选择与 masks 中 True 值对应的预测关键点。 pred_kpt 是这些预测关键点的集合。
pred_kpt = pred_kpts[masks]
# gt_kpt[..., 2] != 0 创建一个掩码,表示每个关键点的可见性。如果关键点的第三个坐标(通常是可见性)不为0,则该关键点被认为是可见的。
# if gt_kpt.shape[-1] == 3 检查 gt_kpt 的最后一个维度是否为3,即是否包含可见性信息。如果包含,则使用上述掩码;否则,创建一个全为 True 的掩码,表示所有关键点都是可见的。
kpt_mask = gt_kpt[..., 2] != 0 if gt_kpt.shape[-1] == 3 else torch.full_like(gt_kpt[..., 0], True)
# 计算 关键点位置损失 。这个损失函数通常考虑了预测关键点与真实关键点之间的差异,并使用 kpt_mask 和 area 来加权这些差异。
kpts_loss = self.keypoint_loss(pred_kpt, gt_kpt, kpt_mask, area) # pose loss
# 检查预测关键点的最后一个维度是否为3,即是否包含可见性信息。
if pred_kpt.shape[-1] == 3:
# 计算关键点对象损失,这是一个二元交叉熵损失,用于分类关键点是否可见。
kpts_obj_loss = self.bce_pose(pred_kpt[..., 2], kpt_mask.float()) # keypoint obj loss
# 返回计算得到的 关键点位置损失 kpts_loss 和 关键点对象损失 kpts_obj_loss 。
return kpts_loss, kpts_obj_loss
# 这段代码的主要目的是计算关键点位置损失和关键点对象损失。检查 masks 中是否有任何 True 值。从 selected_keypoints 中选择与 masks 中 True 值对应的真实关键点。计算这些框的面积。从 pred_kpts 中选择与 masks 中 True 值对应的预测关键点。创建关键点的可见性掩码。计算关键点位置损失。如果预测关键点包含可见性信息,计算关键点对象损失。返回计算得到的损失值。通过这些步骤,可以评估模型在关键点检测任务中的性能,并为反向传播提供必要的梯度信息。
# alculate_keypoints_loss 方法的主要目的是计算预测关键点与真实关键点之间的位置损失和可见性损失。根据批次索引 batch_idx 将关键点分批存储。使用 target_gt_idx 选择与预测锚点对应的真实关键点。将选择的关键点坐标除以步长,转换为绝对坐标。计算关键点位置损失和关键点对象损失。返回计算得到的损失值。这个方法在关键点检测和姿态估计任务中非常重要,因为它能够评估模型的预测精度,并为反向传播提供必要的梯度信息。
# 这段代码定义了一个名为 v8ClassificationLoss 的类,用于计算分类任务的训练损失。
# 定义了一个名为 v8ClassificationLoss 的类,这个类用于封装分类损失的计算逻辑。
class v8ClassificationLoss:
# 计算训练损失的标准类。
"""Criterion class for computing training losses."""
# 定义了一个 __call__ 方法,这个方法使得类的实例可以像函数一样被调用。它接受两个参数。
# 1.preds :预测结果。
# 2.batch :包含真实标签的批次数据。
def __call__(self, preds, batch):
# 计算预测和真实标签之间的分类损失。
"""Compute the classification loss between predictions and true labels."""
# 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' )。这个值可以用于反向传播,以更新模型的权重。
# 使用PyTorch的 torch.nn.functional.cross_entropy 函数计算分类损失。这个函数接受预测结果 preds 和真实标签 batch["cls"] 作为输入,并设置 reduction="mean" 参数,表示将损失值在批次维度上进行平均。计算得到的损失值存储在变量 loss 中。
loss = torch.nn.functional.cross_entropy(preds, batch["cls"], reduction="mean")
# 将计算得到的损失值 loss 从计算图中分离出来,得到一个不参与梯度计算的张量 loss_items 。这通常用于在不干扰梯度传播的情况下,将损失值用于其他目的,如监控训练过程中的损失变化。
loss_items = loss.detach()
# 返回两个值, loss (参与梯度计算的损失值)和 loss_items (分离后的损失值,可用于监控等其他用途)。
return loss, loss_items
# v8ClassificationLoss 类封装了分类损失的计算过程,通过实现 __call__ 方法,使得类的实例可以方便地被调用以计算损失。这个类使用PyTorch的 cross_entropy 函数计算分类损失,并提供了两个返回值:一个用于参与梯度计算,另一个用于其他用途(如监控)。这种设计使得损失计算既简洁又灵活,便于在深度学习训练过程中使用。
# 这段代码定义了一个名为 v8OBBLoss 的类,它继承自 v8DetectionLoss 类,用于计算目标检测任务中的定向边界框(OBB)损失。
# 定义了一个名为 v8OBBLoss 的类,继承自 v8DetectionLoss 类。
class v8OBBLoss(v8DetectionLoss):
# 这段代码是 v8OBBLoss 类的初始化方法 __init__ 的定义,用于初始化类的实例。
# 定义了 v8OBBLoss 类的初始化方法 __init__ ,它接受一个参数。
# 1.model :这个参数通常是一个模型实例,用于传递模型的相关信息和配置。
def __init__(self, model):
# 使用模型、分配器和旋转的 bbox 损失初始化 v8OBBLoss。
# 注意模型必须解并行。
"""
Initializes v8OBBLoss with model, assigner, and rotated bbox loss.
Note model must be de-paralleled.
"""
# 调用父类 v8DetectionLoss 的初始化方法,将 model 参数传递给父类。这一步确保了父类的初始化逻辑得以执行,例如初始化一些在父类中定义的属性或方法。
super().__init__(model)
# 初始化一个 RotatedTaskAlignedAssigner 实例,并将其赋值给类的属性 assigner 。 RotatedTaskAlignedAssigner 是一个用于分配正负样本的类,它根据预测结果和真实标签来确定哪些预测框是正样本(与真实目标匹配的预测框)和哪些是负样本(不与真实目标匹配的预测框)。
# topk=10 :选择每个真实目标的前10个最匹配的预测框。
# num_classes=self.nc :类别数, self.nc 是从模型或父类中继承的属性,表示数据集中的类别总数。
# alpha=0.5 和 beta=6.0 :这两个参数是分配器使用的超参数,用于控制分配正负样本时的策略,具体作用与损失函数的计算或样本权重的分配有关。
self.assigner = RotatedTaskAlignedAssigner(topk=10, num_classes=self.nc, alpha=0.5, beta=6.0)
# 初始化一个 RotatedBboxLoss 实例,并将其赋值给类的属性 bbox_loss 。 RotatedBboxLoss 是一个用于计算旋转边界框损失的类。
# self.reg_max - 1 :最大回归值减1, self.reg_max 是模型中定义的一个属性,表示回归任务中的最大值范围。
# use_dfl=self.use_dfl :是否使用分布焦点损失(DFL), self.use_dfl 是从模型或父类中继承的属性,表示是否启用DFL来计算损失。
# .to(self.device) :将 RotatedBboxLoss 实例移动到指定的设备(如CPU或GPU), self.device 是从模型或父类中继承的属性,表示当前使用的计算设备。
self.bbox_loss = RotatedBboxLoss(self.reg_max - 1, use_dfl=self.use_dfl).to(self.device)
# 这个初始化方法的主要作用是设置 v8OBBLoss 类的两个核心组件,样本分配器 assigner 和边界框损失计算模块 bbox_loss 。通过这些组件的初始化, v8OBBLoss 类能够有效地计算定向边界框(OBB)任务中的损失,包括分类损失和回归损失,这对于训练一个有效的目标检测模型至关重要。
# 这段代码定义了 v8OBBLoss 类中的 preprocess 方法,其作用是对目标数据进行预处理,以便后续计算损失。
# 定义了 preprocess 方法,接受三个参数。
# 1.targets :目标数据,形状为 [n, 6] ,其中 n 是目标的数量,每一行包含 [image_index, class_id, x1, y1, x2, y2] 。
# 2.batch_size :批次大小。
# 3.scale_tensor :用于缩放边界框坐标的张量,形状为 [4] ,通常用于将边界框坐标从模型的输出尺度转换到图像的实际尺度。
def preprocess(self, targets, batch_size, scale_tensor):
# 预处理目标计数并与输入批次大小匹配以输出张量。
"""Preprocesses the target counts and matches with the input batch size to output a tensor."""
# 如果目标数据为空(即没有目标),则返回一个形状为 [batch_size, 0, 6] 的全零张量。这个张量表示没有目标的批次,每个图像都没有目标。
if targets.shape[0] == 0:
out = torch.zeros(batch_size, 0, 6, device=self.device)
else:
# 提取目标数据中的图像索引, i 的形状为 [n] ,表示每个目标所属的图像索引。
i = targets[:, 0] # image index
# 使用 unique 方法获取每个图像索引的 唯一值 及 其出现次数 , counts 的形状为 [batch_size] ,表示每个图像中的目标数量。
_, counts = i.unique(return_counts=True)
# 将 counts 转换为 int32 类型。
counts = counts.to(dtype=torch.int32)
# 创建一个形状为 [batch_size, max(counts), 6] 的全零张量 out ,用于 存储预处理后的目标数据 。 max(counts) 是批次中目标数量最多的图像的目标数量。
out = torch.zeros(batch_size, counts.max(), 6, device=self.device)
# 遍历每个图像索引。
for j in range(batch_size):
# 找到属于第 j 个图像的所有目标, matches 是一个布尔张量,形状为 [n] ,表示哪些目标属于第 j 个图像。
matches = i == j
# 计算第 j 个图像中的目标数量。
n = matches.sum()
# 如果第 j 个图像有目标,则提取这些目标的边界框坐标, bboxes 的形状为 [n, 4] ,表示这些目标的边界框坐标。
if n:
bboxes = targets[matches, 2:]
# 将边界框坐标乘以 scale_tensor ,进行尺度转换。
bboxes[..., :4].mul_(scale_tensor)
# 将 类别标签 和 转换后的边界框坐标 拼接在一起,存储到 out 张量中。 targets[matches, 1:2] 提取类别标签,形状为 [n, 1] , bboxes 的形状为 [n, 4] ,拼接后的形状为 [n, 5] 。
out[j, :n] = torch.cat([targets[matches, 1:2], bboxes], dim=-1)
# 返回预处理后的目标数据 out ,形状为 [batch_size, max(counts), 6] 。
return out
# preprocess 方法的主要作用是将目标数据从原始格式转换为适合后续损失计算的格式。检查目标数据是否为空,如果为空则返回全零张量。提取每个图像的目标数量,并创建一个足够大的张量来存储预处理后的目标数据。遍历每个图像,提取属于该图像的目标,将边界框坐标进行尺度转换,并将类别标签和转换后的边界框坐标拼接在一起。返回预处理后的目标数据。这个方法确保了目标数据的格式统一,便于后续的损失计算。
# 这段代码定义了 v8OBBLoss 类中的 __call__ 方法,使得类的实例可以像函数一样被调用,用于计算定向边界框(OBB)任务的总损失。
# 定义了 __call__ 方法,接受两个参数。
# 1.preds :模型的预测结果。
# 2.batch :当前批次的数据,包含真实标签和边界框信息。
def __call__(self, preds, batch):
# 计算并返回 YOLO 模型的损失。
"""Calculate and return the loss for the YOLO model."""
# 这段代码是 v8OBBLoss 类中的 __call__ 方法的一部分,用于初始化损失值并提取预测结果中的关键信息。
# 初始化一个长度为3的零张量 loss ,用于存储三种损失值。 loss[0] :边界框损失(box loss)。 loss[1] :分类损失(classification loss)。 loss[2] :分布焦点损失(distribution focal loss, DFL)。
loss = torch.zeros(3, device=self.device) # box, cls, dfl
# 从预测结果 preds 中提取 特征图 feats 和 预测角度 pred_angle 。 preds 的结构是一个列表或一个包含多个元素的元组。如果 preds[0] 是一个列表,直接提取;否则,提取 preds[1] 。这一步确保了代码的灵活性,可以处理不同格式的预测结果。
feats, pred_angle = preds if isinstance(preds[0], list) else preds[1]
# 获取 批次大小 batch_size ,即 预测角度 pred_angle 的第一个维度的大小。这表示当前批次中有多少个样本。
batch_size = pred_angle.shape[0] # batch size, number of masks, mask height, mask width
# 这行代码的目的是从特征图 feats 中提取 预测分布 pred_distri 和 预测分数 pred_scores 。具体步骤如下 :
# [xi.view(feats[0].shape[0], self.no, -1) for xi in feats] :对每个特征图 xi 进行展平操作,将其形状从 [batch_size, channels, height, width] 转换为 [batch_size, self.no, -1] 。 self.no 是每个锚点的输出数量, -1 表示自动计算最后一个维度的大小。
# torch.cat([xi.view(feats[0].shape[0], self.no, -1) for xi in feats], 2) :将所有展平后的特征图在最后一个维度上进行拼接,得到一个形状为 [batch_size, self.no, total_anchors] 的张量。 total_anchors 表示在当前特征图上所有可能的锚点(anchors)的总数。
# .split((self.reg_max * 4, self.nc), 1) :将拼接后的张量在第二个维度上分割为两部分 :
# pred_distri : 预测分布 ,形状为 [batch_size, total_anchors, self.reg_max * 4] ,表示每个锚点的四个边界框坐标的预测分布。
# pred_scores :预测分数,形状为 [batch_size, total_anchors, 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
)
# 这段代码的主要作用是初始化损失值,并从预测结果中提取关键信息,包括特征图、预测角度、预测分布和预测分数。这些信息将用于后续的损失计算,确保模型的训练过程能够正确地评估预测结果与真实标签之间的差异。
# 这段代码继续处理预测结果,调整张量的维度顺序,并计算图像大小和锚点坐标。
# b, grids, ..
# 调整 预测分数 pred_scores 、 预测分布 pred_distri 和 预测角度 pred_angle 的维度顺序,使其符合后续计算的要求。具体步骤如下:
# .permute(0, 2, 1) :将张量的维度从 [batch_size, self.no, total_anchors] 调整为 [batch_size, total_anchors, self.no] 。这一步确保了每个锚点的预测结果在第二个维度上。
# .contiguous() :调用 contiguous() 方法确保张量在内存中是连续的。这一步在某些操作(如视图操作)中是必要的,因为 permute 操作可能会导致张量在内存中变得不连续。
pred_scores = pred_scores.permute(0, 2, 1).contiguous()
pred_distri = pred_distri.permute(0, 2, 1).contiguous()
pred_angle = pred_angle.permute(0, 2, 1).contiguous()
# 获取预测分数 pred_scores 的数据类型。
dtype = pred_scores.dtype
# feats[0].shape[2:] 获取第一个特征图的 高度 和 宽度 。 torch.tensor(..., device=self.device, dtype=dtype) 将高度和宽度转换为张量,并确保其在指定的设备上且数据类型与 pred_scores 相同。 * self.stride[0] 将特征图的尺寸乘以步长 self.stride[0] ,得到实际图像的尺寸。
imgsz = torch.tensor(feats[0].shape[2:], device=self.device, dtype=dtype) * self.stride[0] # image size (h,w)
# 调用 make_anchors 函数,生成锚点坐标和步长张量。
# 生成锚点坐标 : feats 特征图列表。 self.stride 步长列表,表示每个特征图的步长。 0.5 偏移量,用于调整锚点的位置。
# 返回值 : anchor_points 生成的锚点坐标,形状为 [total_anchors, 2] 。 stride_tensor 每个锚点对应的步长,形状为 [total_anchors, 1] 。
anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5)
# 这段代码调整预测分数、预测分布和预测角度的维度顺序,确保每个锚点的预测结果在第二个维度上。计算图像的大小,确保其数据类型与预测分数相同。生成锚点坐标和步长张量,用于后续的边界框解码和损失计算。这些步骤为后续的损失计算提供了必要的准备,确保预测结果和真实标签之间的差异可以正确评估。
# 这段代码处理目标数据,将其从原始格式转换为适合后续损失计算的格式。具体步骤包括提取目标索引、类别标签和边界框坐标,过滤掉过小的边界框,调用 preprocess 方法进行预处理,并分离出目标标签和边界框。
# targets
try:
# 从批次数据 batch 中提取 图像索引 batch_idx ,并将其形状调整为 [-1, 1] ,表示 每个目标所属的图像索引 。
batch_idx = batch["batch_idx"].view(-1, 1)
# 将 图像索引 、 类别标签 batch["cls"] 和 边界框坐标 batch["bboxes"] 拼接在一起,形成一个形状为 [-1, 7] 的张量 targets 。每个目标的表示形式为 [image_index, class_id, x1, y1, x2, y2, angle] 。
targets = torch.cat((batch_idx, batch["cls"].view(-1, 1), batch["bboxes"].view(-1, 5)), 1)
# 计算每个目标的 宽度 rw 和 高度 rh ,并将其从特征图尺度转换为图像实际尺度。 imgsz[0].item() 和 imgsz[1].item() 分别表示图像的高度和宽度。
rw, rh = targets[:, 4] * imgsz[0].item(), targets[:, 5] * imgsz[1].item()
# 过滤掉宽度和高度小于2的边界框,以稳定训练过程。这一步确保了模型不会被过小的边界框影响,提高训练的稳定性和效果。
targets = targets[(rw >= 2) & (rh >= 2)] # filter rboxes of tiny size to stabilize training
# 调用 preprocess 方法对目标数据进行预处理。 targets.to(self.device) 将目标数据移动到指定的设备(如GPU), batch_size 表示批次大小, scale_tensor=imgsz[[1, 0, 1, 0]] 用于将边界框坐标从特征图尺度转换为图像实际尺度。
targets = self.preprocess(targets.to(self.device), batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
# 将 预处理后的目标数据 targets 分离为 目标标签 gt_labels 和 目标边界框 gt_bboxes 。 gt_labels 的形状为 [batch_size, max_targets, 1] , gt_bboxes 的形状为 [batch_size, max_targets, 5] ,其中 max_targets 是每个图像中目标数量的最大值。
gt_labels, gt_bboxes = targets.split((1, 5), 2) # cls, xywhr
# 生成一个 前景掩码 mask_gt ,用于标记哪些目标是有效的。 gt_bboxes.sum(2, keepdim=True) 计算每个目标边界框的坐标和角度的和, gt_(0) 生成一个布尔张量,表示哪些目标的和大于0,即有效的目标。
mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0)
# 如果在处理目标数据时发生 RuntimeError ,则抛出一个 TypeError ,提示数据集格式不正确或不是OBB数据集。这一步确保了数据集的格式正确,避免在训练过程中出现意外错误。
except RuntimeError as e:
raise TypeError(
"ERROR ❌ OBB dataset incorrectly formatted or not a OBB dataset.\n" # 错误 ❌ OBB 数据集格式不正确或不是 OBB 数据集。
"This error can occur when incorrectly training a 'OBB' model on a 'detect' dataset, " # 当在“检测”数据集上错误地训练“OBB”模型时,可能会发生此错误,
"i.e. 'yolo train model=yolov8n-obb.pt data=dota8.yaml'.\nVerify your dataset is a " # 即“yolo train model=yolov8n-obb.pt data=dota8.yaml”。
"correctly formatted 'OBB' dataset using 'data=dota8.yaml' " # n请使用“data=dota8.yaml”作为示例,验证您的数据集是否为格式正确的“OBB”数据集。
"as an example.\nSee https://docs.ultralytics.com/datasets/obb/ for help." # n有关帮助,请参阅 https://docs.ultralytics.com/datasets/obb/。
) from e
# 这段代码的主要作用是处理目标数据,确保其格式正确,并过滤掉过小的边界框,以提高训练的稳定性和效果。通过这些步骤,目标数据被转换为适合后续损失计算的格式,确保模型能够正确地评估预测结果与真实标签之间的差异。
# 这段代码处理预测的边界框,并使用分配器(assigner)将预测的边界框与真实目标进行匹配,生成用于损失计算的目标边界框、目标分数和前景掩码。
# Pboxes
# 调用 bbox_decode 方法,将 预测的分布 pred_distri 、 预测的角度 pred_angle 和 锚点坐标 anchor_points 解码为 预测的边界框 pred_bboxes 。解码后的边界框格式为 xyxy ,形状为 [batch_size, h*w, 4] 。
pred_bboxes = self.bbox_decode(anchor_points, pred_distri, pred_angle) # xyxy, (b, h*w, 4)
# 克隆并分离 预测的边界框 pred_bboxes ,生成 bboxes_for_assigner ,用于 分配器的输入 。这一步确保在分配器中对边界框进行操作时不会影响原始的预测边界框。
bboxes_for_assigner = pred_bboxes.clone().detach()
# Only the first four elements need to be scaled
# 将 预测的边界框 bboxes_for_assigner 的前四个元素(即边界框的坐标)乘以步长张量 stride_tensor ,将其从特征图尺度转换为图像实际尺度。
bboxes_for_assigner[..., :4] *= stride_tensor
# 调用分配器 self.assigner ,将预测的分数、预测的边界框、锚点坐标、目标标签、目标边界框和前景掩码进行匹配,生成 用于损失计算的目标边界框 target_bboxes 、 目标分数 target_scores 和 前景掩码 fg_mask 。
_, target_bboxes, target_scores, fg_mask, _ = self.assigner(
# 预测的分数 经过sigmoid激活函数,形状为 [batch_size, h*w, num_classes] 。
pred_scores.detach().sigmoid(),
# 预测的边界框,数据类型与目标边界框 gt_bboxes 相同。
bboxes_for_assigner.type(gt_bboxes.dtype),
# 锚点坐标乘以步长张量,形状为 [h*w, 2] 。
anchor_points * stride_tensor,
# 目标标签,形状为 [batch_size, max_targets, 1] 。
gt_labels,
# 目标边界框,形状为 [batch_size, max_targets, 5]
gt_bboxes,
# 前景掩码,形状为 [batch_size, max_targets, 1] 。
mask_gt,
)
# 这段代码解码预测的边界框,将其从特征图尺度转换为图像实际尺度。准备分配器的输入,包括预测的分数、预测的边界框和锚点坐标。使用分配器将预测的边界框与真实目标进行匹配,生成用于损失计算的目标边界框、目标分数和前景掩码。这些步骤确保了预测结果与真实标签之间的匹配关系被正确地建立,为后续的损失计算提供了必要的输入。通过这些步骤,模型能够有效地评估预测结果与真实标签之间的差异,从而优化模型的性能。
# 这段代码计算分类损失和边界框损失,并应用超参数增益,最后返回总损失和分离后的损失值。
# 计算目标分数 target_scores 的总和,并确保其不为零,以避免除以零的错误。如果总和为零,则将其设置为1。
target_scores_sum = max(target_scores.sum(), 1)
# Cls loss
# loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum # VFL way
# 计算分类损失,使用二元交叉熵损失(BCE)。
# self.bce(pred_scores, target_scores.to(dtype)) :计算 预测分数 pred_scores 和 目标分数 target_scores 之间的二元交叉熵损失。
# .sum() :对所有元素求和。
# / target_scores_sum :将总和除以目标分数的总和,以归一化损失值。
loss[1] = self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum # BCE
# Bbox loss
# 如果存在前景样本(即 fg_mask.sum() > 0 ),则计算边界框损失和分布焦点损失(DFL)。
if fg_mask.sum():
# 将目标边界框的坐标除以步长张量,将其从图像实际尺度转换为特征图尺度。
target_bboxes[..., :4] /= stride_tensor
# 调用 bbox_loss 方法,计算 边界框损失 和 分布焦点损失 。 loss[0] 边界框损失。 loss[2] 分布焦点损失。
loss[0], loss[2] = self.bbox_loss(
pred_distri, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask
)
# 如果不存在前景样本,则将 预测角度的零和 添加到边界框损失中,以避免损失值为零。
else:
loss[0] += (pred_angle * 0).sum()
# 将损失值乘以相应的超参数增益,以调整损失值的权重。
# self.hyp.box 边界框损失的增益。
loss[0] *= self.hyp.box # box gain
# self.hyp.cls 分类损失的增益。
loss[1] *= self.hyp.cls # cls gain
# self.hyp.dfl 分布焦点损失的增益。
loss[2] *= self.hyp.dfl # dfl gain
# 返回 总损失 和 分离后的损失值 。 loss.sum() * batch_size 计算总损失,并乘以批次大小,以得到批次的总损失。 loss.detach() 分离损失值,避免在后续计算中影响梯度传播。
return loss.sum() * batch_size, loss.detach() # loss(box, cls, dfl)
# 这段代码计算目标分数的总和,确保其不为零。计算分类损失,使用二元交叉熵损失(BCE)。计算边界框损失和分布焦点损失,如果存在前景样本。应用超参数增益,调整损失值的权重。返回总损失和分离后的损失值。这些步骤确保了损失计算的完整性和准确性,是训练定向边界框检测模型的关键部分。通过这些步骤,模型能够有效地评估预测结果与真实标签之间的差异,从而优化模型的性能。
# __call__ 方法的主要作用是计算定向边界框(OBB)任务的总损失,包括边界框损失、分类损失和分布焦点损失。初始化损失值。提取预测结果中的特征图、预测分布、预测分数和预测角度。调整预测结果的维度顺序。计算图像大小和锚点坐标。处理目标数据,包括提取目标索引、类别标签和边界框坐标,过滤掉过小的边界框,预处理目标数据。解码预测的边界框坐标。准备分配器输入,调用分配器分配正负样本。计算分类损失和边界框损失。应用超参数增益。 返回总损失和分离后的损失值。这个方法确保了损失计算的完整性和准确性,是训练定向边界框检测模型的关键部分。
# 这段代码定义了 v8OBBLoss 类中的 bbox_decode 方法,用于将预测的边界框分布和角度解码为实际的旋转边界框坐标。
# 定义了 bbox_decode 方法,接受三个参数。
# 1.anchor_points :锚点坐标,形状为 [h*w, 2] ,其中 h 和 w 分别是特征图的高度和宽度。
# 2.pred_dist :预测的旋转距离,形状为 [bs, h*w, 4] ,其中 bs 是批次大小, 4 表示边界框的四个坐标点。
# 3.pred_angle :预测的角度,形状为 [bs, h*w, 1] 。
def bbox_decode(self, anchor_points, pred_dist, pred_angle):
# 从锚点和分布解码预测的对象边界框坐标。
# 参数:
# anchor_points (torch.Tensor):锚点,(h*w, 2)。
# pred_dist (torch.Tensor):预测的旋转距离,(bs, h*w, 4)。
# pred_angle (torch.Tensor):预测的角度,(bs, h*w, 1)。
# 返回:
# (torch.Tensor):预测的旋转边界框和角度,(bs, h*w, 5)。
"""
Decode predicted object bounding box coordinates from anchor points and distribution.
Args:
anchor_points (torch.Tensor): Anchor points, (h*w, 2).
pred_dist (torch.Tensor): Predicted rotated distance, (bs, h*w, 4).
pred_angle (torch.Tensor): Predicted angle, (bs, h*w, 1).
Returns:
(torch.Tensor): Predicted rotated bounding boxes with angles, (bs, h*w, 5).
"""
# 判断是否启用了分布焦点损失(DFL)。如果启用了DFL,执行以下步骤
if self.use_dfl:
# 取预测分布 pred_dist 的形状,并将其分解为三个维度。 b 批次大小(batch size)。 a 锚点数量(anchors),即特征图上的每个位置。 c 通道数(channels),表示预测的分布的维度。
b, a, c = pred_dist.shape # batch, anchors, channels
# 对预测分布进行处理,具体步骤如下 :
# pred_dist.view(b, a, 4, c // 4) :将 pred_dist 的形状从 [b, a, c] 重新调整为 [b, a, 4, c // 4] 。这里假设 c 是4的倍数, c // 4 表示每个边界框的四个坐标点(x1, y1, x2, y2)的分布维度。
# .softmax(3) :对最后一个维度(即分布维度)应用softmax函数。softmax函数将每个分布转换为概率分布,使得每个值在0到1之间,并且所有值的和为1。这一步确保了每个分布是一个有效的概率分布。
# .matmul(self.proj.type(pred_dist.dtype)) :使用矩阵乘法将概率分布转换为实际的距离值。 self.proj 是一个预定义的投影矩阵,其形状为 [c // 4, 1] 。 self.proj.type(pred_dist.dtype) 确保投影矩阵的数据类型与 pred_dist 相同,以避免类型不匹配的问题。矩阵乘法的结果是一个形状为 [b, a, 4, 1] 的张量,表示每个边界框的四个坐标点的实际距离值。
pred_dist = pred_dist.view(b, a, 4, c // 4).softmax(3).matmul(self.proj.type(pred_dist.dtype))
# 将解码后的 边界框坐标 和 预测的角度 拼接在一起,返回最终的旋转边界框坐标。
# 调用 dist2rbox 函数,将预测的分布 pred_dist 、预测的角度 pred_angle 和锚点坐标 anchor_points 转换为实际的旋转边界框坐标,形状为 [bs, h*w, 4] 。 将预测的角度 pred_angle 拼接到解码后的边界框坐标后面,最终形状为 [bs, h*w, 5] 。
# def dist2rbox(pred_dist, pred_angle, anchor_points, dim=-1):
# -> 用于将模型输出的距离预测( pred_dist )、角度预测( pred_angle )和锚点( anchor_points )转换为旋转边界框(Oriented Bounding Boxes, OBB)的坐标。将边界框的中心坐标和宽度高度( lt + rb )连接起来,形成完整的旋转边界框坐标,并返回。
# -> return torch.cat([xy, lt + rb], dim=dim)
return torch.cat((dist2rbox(pred_dist, pred_angle, anchor_points), pred_angle), dim=-1)
# bbox_decode 方法的主要作用是将预测的分布和角度解码为实际的旋转边界框坐标。如果启用了DFL,对预测的分布进行softmax处理,并使用投影矩阵将其转换为实际的距离值。调用 dist2rbox 函数,将预测的分布、预测的角度和锚点坐标转换为实际的旋转边界框坐标。将预测的角度拼接到解码后的边界框坐标后面,返回最终的旋转边界框坐标。这个方法确保了预测的边界框和角度可以正确地转换为实际的旋转边界框坐标,便于后续的损失计算。
# 这段代码定义了一个名为 v10DetectLoss 的类,用于计算目标检测任务中的两种不同策略(one-to-many和one-to-one)的损失,并将它们组合起来。
# 定义了一个名为 v10DetectLoss 的类,用于计算目标检测任务中的损失。
class v10DetectLoss:
# 定义了类的初始化方法 __init__ ,接受一个参数。
# 1.model :表示模型实例。
def __init__(self, model):
# 在初始化方法中,创建了两个 v8DetectionLoss 实例。
# 用于计算one-to-many策略的损失, tal_topk=10 表示每个真实目标选择前10个最匹配的预测框。
self.one2many = v8DetectionLoss(model, tal_topk=10)
# 用于计算one-to-one策略的损失, tal_topk=1 表示每个真实目标选择1个最匹配的预测框。
self.one2one = v8DetectionLoss(model, tal_topk=1)
# 定义了 __call__ 方法,使得类的实例可以像函数一样被调用。接受两个参数。
# 1.preds :模型的预测结果,包含 one2many 和 one2one 两个部分。
# 2.batch :当前批次的数据,包含真实标签和边界框信息。
def __call__(self, preds, batch):
# 从预测结果 preds 中提取 one2many 部分的预测结果,并调用 self.one2many 计算one-to-many策略的损失。 loss_one2many 是一个元组,包含 总损失 和 分离后的损失值 。
one2many = preds["one2many"]
loss_one2many = self.one2many(one2many, batch)
# 从预测结果 preds 中提取 one2one 部分的预测结果,并调用 self.one2one 计算one-to-one策略的损失。 loss_one2one 是一个元组,包含 总损失 和 分离后的损失值 。
one2one = preds["one2one"]
loss_one2one = self.one2one(one2one, batch)
# 返回两种策略的总损失之和,以及将两种策略的分离后的损失值拼接在一起的张量。
return loss_one2many[0] + loss_one2one[0], torch.cat((loss_one2many[1], loss_one2one[1]))
# v10DetectLoss 类初始化两个 v8DetectionLoss 实例,分别用于计算one-to-many和one-to-one策略的损失。从预测结果中提取 one2many 和 one2one 部分的预测结果。分别计算两种策略的损失。返回两种策略的总损失之和,以及将两种策略的分离后的损失值拼接在一起的张量。这种设计允许模型在训练过程中同时考虑两种不同的匹配策略,从而提高模型的鲁棒性和性能。通过组合两种策略的损失,模型可以更好地平衡不同情况下的预测结果。