YOLOv10-1.1部分代码阅读笔记-ops.py
ops.py
ultralytics\models\utils\ops.py
目录
ops.py
1.所需的库和模块
2.class HungarianMatcher(nn.Module):
3.def get_cdn_group(batch, num_classes, num_queries, class_embed, num_dn=100, cls_noise_ratio=0.5, box_noise_scale=1.0, training=False):
1.所需的库和模块
# Ultralytics YOLO 🚀, AGPL-3.0 license
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment
from ultralytics.utils.metrics import bbox_iou
from ultralytics.utils.ops import xywh2xyxy, xyxy2xywh
2.class HungarianMatcher(nn.Module):
# 匈牙利算法是一种解决二分图最大匹配问题的经典算法。下面详细介绍其概念、原理及应用。
# 概念 :
# 在图论中,二分图是指图的顶点可以分为两个互不相交的集合,使得图中的每条边都连接一个集合中的顶点到另一个集合中的顶点。最大匹配是指在二分图中选取尽可能多的边,使得这些边没有公共的顶点。匈牙利算法就是用来求解这种二分图的最大匹配问题的。
# 原理 :
# 增广路径 :匈牙利算法的核心思想是寻找增广路径。增广路径是一条交替经过未匹配边和匹配边的路径,且起点和终点都是未匹配的顶点。如果找到了一条增广路径,就可以通过改变这条路径上边的匹配状态(未匹配边变为匹配边,匹配边变为未匹配边)来增加匹配的数量。
# 算法步骤 :
# 初始化匹配集合为空。
# 对于二分图中的每个左部顶点,尝试为其寻找匹配 :
# 从该左部顶点开始进行深度优先搜索(DFS)或广度优先搜索(BFS),寻找增广路径。
# 如果找到增广路径,则通过改变增广路径上边的匹配状态来增加匹配,并将该左部顶点标记为已匹配。
# 如果未找到增广路径,则继续尝试下一个左部顶点。
# 重复上述过程,直到所有左部顶点都尝试过或无法再找到增广路径为止。
# 应用 :
# 目标匹配 :在目标检测任务中,匈牙利算法可以用于将预测的目标与真实的目标进行匹配。例如在 HungarianMatcher 类中,通过构建成本矩阵,将预测目标和真实目标之间的匹配问题转化为二分图匹配问题,然后利用匈牙利算法找到最优匹配,从而为后续的损失计算等操作提供依据。
# 任务分配 :假设有一些任务和一些工人,每个工人完成每个任务都有一定的成本。匈牙利算法可以用来将任务分配给工人,使得总成本最小。将任务和工人分别作为二分图的两个集合,成本作为边的权重,通过匈牙利算法找到最优的分配方案。
# 网络流问题 :匈牙利算法也可以应用于某些网络流问题。例如,在一个网络中,可以将节点分为两组,一组表示源点到中间节点的连接,另一组表示中间节点到汇点的连接,通过构建二分图并使用匈牙利算法求解最大匹配,从而得到网络的最大流。
# 匈牙利算法因其高效性和实用性,在计算机科学、运筹学等领域得到了广泛的应用。
# 这段代码定义了一个名为 HungarianMatcher 的类,继承自 nn.Module ,用于实现基于匈牙利算法的目标检测中的匹配过程。
# 定义了一个名为 HungarianMatcher 的类,继承自 PyTorch 的 nn.Module 基类,这样可以方便地与其他 PyTorch 模块进行集成和使用。
class HungarianMatcher(nn.Module):
# 实现 HungarianMatcher 的模块,这是一个可微分模块,用于以端到端方式解决分配问题。
# HungarianMatcher 使用考虑分类分数、边界框坐标以及可选的掩码预测的成本函数对预测和地面真实边界框执行最佳分配。
# 方法:
# forward(pred_bboxes, pred_scores, gt_bboxes, gt_cls, gt_groups, mask=None, gt_mask=None): 计算一批预测结果和基本事实之间的分配。
# _cost_mask(bs, num_gts, mask=None, gt_mask=None): 如果预测了 mask,则计算 mask 成本和 dice 成本。
"""
A module implementing the HungarianMatcher, which is a differentiable module to solve the assignment problem in an
end-to-end fashion.
HungarianMatcher performs optimal assignment over the predicted and ground truth bounding boxes using a cost
function that considers classification scores, bounding box coordinates, and optionally, mask predictions.
Attributes:
cost_gain (dict): Dictionary of cost coefficients: 'class', 'bbox', 'giou', 'mask', and 'dice'.
use_fl (bool): Indicates whether to use Focal Loss for the classification cost calculation.
with_mask (bool): Indicates whether the model makes mask predictions.
num_sample_points (int): The number of sample points used in mask cost calculation.
alpha (float): The alpha factor in Focal Loss calculation.
gamma (float): The gamma factor in Focal Loss calculation.
Methods:
forward(pred_bboxes, pred_scores, gt_bboxes, gt_cls, gt_groups, masks=None, gt_mask=None): Computes the
assignment between predictions and ground truths for a batch.
_cost_mask(bs, num_gts, masks=None, gt_mask=None): Computes the mask cost and dice cost if masks are predicted.
"""
# 这段代码定义了 HungarianMatcher 类的初始化方法 __init__ ,用于设置该类实例的初始状态和参数。
# 定义了 HungarianMatcher 类的初始化方法,它接受以下参数 :
# 1.cost_gain :一个字典,用于指定不同成本项的权重,默认为 None 。
# 2.use_fl :一个布尔值,表示是否使用Focal Loss,默认为 True 。
# 3.with_mask :一个布尔值,表示是否考虑掩码信息,默认为 False 。
# 4.num_sample_points :一个整数,表示采样点的数量,默认为 12544 。
# 5.alpha 和 6.gamma :两个浮点数,用于Focal Loss的计算,默认分别为 0.25 和 2.0 。
def __init__(self, cost_gain=None, use_fl=True, with_mask=False, num_sample_points=12544, alpha=0.25, gamma=2.0):
# 使用成本系数、焦点损失、掩模预测、样本点和 alpha gamma 因子初始化 HungarianMatcher。
"""Initializes HungarianMatcher with cost coefficients, Focal Loss, mask prediction, sample points, and alpha
gamma factors.
"""
# 调用父类 nn.Module 的初始化方法,完成基本的初始化操作,确保 HungarianMatcher 类能够正确地作为 PyTorch 模块使用。
super().__init__()
# 如果传入的 cost_gain 为 None ,则使用默认的成本权重字典。这个字典定义了不同成本项的权重,例如分类成本权重为1,边界框成本权重为5,GIoU成本权重为2等。
if cost_gain is None:
cost_gain = {"class": 1, "bbox": 5, "giou": 2, "mask": 1, "dice": 1}
# 将传入的成本权重字典赋值给类的 cost_gain 属性,用于后续成本计算时的权重分配。
self.cost_gain = cost_gain
# 将 use_fl 参数赋值给类的 use_fl 属性,用于后续判断是否使用Focal Loss。Focal Loss是一种改进的交叉熵损失函数,适用于类别不平衡问题。
self.use_fl = use_fl
# 将 with_mask 参数赋值给类的 with_mask 属性,用于后续判断是否考虑掩码信息。在一些任务中,如实例分割,掩码信息是非常重要的。
self.with_mask = with_mask
# 将 num_sample_points 参数赋值给类的 num_sample_points 属性,用于后续可能的采样操作。例如,在处理掩码时,可能需要对掩码进行采样以减少计算量。
self.num_sample_points = num_sample_points
# 将 alpha 参数赋值给类的 alpha 属性,用于Focal Loss的计算。 alpha 是一个超参数,用于调整正负样本的权重。
self.alpha = alpha
# 将 gamma 参数赋值给类的 gamma 属性,用于Focal Loss的计算。 gamma 是一个超参数,用于调整损失函数的聚焦程度。
self.gamma = gamma
# 这段初始化方法 __init__ 的主要作用是设置 HungarianMatcher 类实例的初始参数和状态。这些参数在后续的前向传播方法 forward 中会被使用,用于计算预测目标和真实目标之间的匹配成本,并通过匈牙利算法找到最优匹配。 cost_gain 用于定义不同成本项的权重。 use_fl 用于决定是否使用Focal Loss。 with_mask 用于决定是否考虑掩码信息。 num_sample_points 用于定义采样点的数量。 alpha 和 gamma 用于Focal Loss的计算。通过这些参数的设置, HungarianMatcher 类可以灵活地应用于不同的目标检测任务,支持多种成本计算方式和损失函数。
# 这段代码定义了 HungarianMatcher 类的前向传播方法 forward ,用于计算预测目标和真实目标之间的匹配成本,并通过匈牙利算法找到最优匹配。
# 定义了前向传播方法 forward ,接受以下参数 :
# 1.pred_bboxes :预测的边界框,形状为 (batch_size, num_queries, 4) 。
# 2.pred_scores :预测的得分,形状为 (batch_size, num_queries, num_classes) 。
# 3.gt_bboxes :真实的边界框,形状为 (num_gt, 4) 。
# 4.gt_cls :真实的类别,形状为 (num_gt,) 。
# 5.gt_groups :真实的分组信息,形状为 (batch_size,) 。
# 6.masks 和 7.gt_mask :可选的掩码信息。
def forward(self, pred_bboxes, pred_scores, gt_bboxes, gt_cls, gt_groups, masks=None, gt_mask=None):
# HungarianMatcher 的前向传递。此函数根据预测和基本事实(分类成本、框之间的 L1 成本和框之间的 GIoU 成本)计算成本,并根据这些成本找到预测和基本事实之间的最佳匹配。
# 参数:
# pred_bboxes (Tensor):形状为 [batch_size, num_queries, 4] 的预测边界框。
# pred_scores (Tensor):形状为 [batch_size, num_queries, num_classes] 的预测分数。
# gt_cls (torch.Tensor):形状为 [num_gts, ] 的基本事实类。
# gt_bboxes (torch.Tensor):形状为 [num_gts, 4] 的基本事实边界框。
# gt_groups (List[int]):长度等于批次大小的列表,包含每个图像的基本事实数量。
# mask (Tensor,可选):形状为 [batch_size, num_queries, height, width]。默认为 None。
# gt_mask(List[Tensor],可选):地面实况掩码列表,每个掩码的形状为 [num_masks, Height, Width]。默认为 None。
# 返回:
# (List[Tuple[Tensor, Tensor]]):大小为 batch_size 的列表,每个元素都是一个元组 (index_i, index_j),其中:
# - index_i 是所选预测的索引张量(按顺序)
# - index_j 是相应所选地面实况目标的索引张量(按顺序)
# 对于每个批次元素,它包含:
# len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
"""
Forward pass for HungarianMatcher. This function computes costs based on prediction and ground truth
(classification cost, L1 cost between boxes and GIoU cost between boxes) and finds the optimal matching between
predictions and ground truth based on these costs.
Args:
pred_bboxes (Tensor): Predicted bounding boxes with shape [batch_size, num_queries, 4].
pred_scores (Tensor): Predicted scores with shape [batch_size, num_queries, num_classes].
gt_cls (torch.Tensor): Ground truth classes with shape [num_gts, ].
gt_bboxes (torch.Tensor): Ground truth bounding boxes with shape [num_gts, 4].
gt_groups (List[int]): List of length equal to batch size, containing the number of ground truths for
each image.
masks (Tensor, optional): Predicted masks with shape [batch_size, num_queries, height, width].
Defaults to None.
gt_mask (List[Tensor], optional): List of ground truth masks, each with shape [num_masks, Height, Width].
Defaults to None.
Returns:
(List[Tuple[Tensor, Tensor]]): A list of size batch_size, each element is a tuple (index_i, index_j), where:
- index_i is the tensor of indices of the selected predictions (in order)
- index_j is the tensor of indices of the corresponding selected ground truth targets (in order)
For each batch element, it holds:
len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
"""
# 获取预测得分的 批次大小 bs 、 查询数量 nq 和 类别数量 nc 。
bs, nq, nc = pred_scores.shape
# 如果 所有真实分组 的和为0,说明没有有效的匹配目标,直接返回一个空的匹配结果列表。
if sum(gt_groups) == 0:
return [(torch.tensor([], dtype=torch.long), torch.tensor([], dtype=torch.long)) for _ in range(bs)]
# 这段代码是 HungarianMatcher 类的 forward 方法中的一部分,主要目的是对预测的得分和边界框进行预处理,以便后续计算成本矩阵。
# We flatten to compute the cost matrices in a batch 我们平铺计算一批成本矩阵。
# [batch_size * num_queries, num_classes]
# 将预测的得分 pred_scores 从形状 (batch_size, num_queries, num_classes) 展平为形状 (batch_size * num_queries, num_classes) 。这样做的目的是为了在批量中计算成本矩阵时,能够更方便地处理数据。
# detach() :调用 detach() 方法,使 pred_scores 不再参与梯度计算。这在计算成本时很有用,因为成本计算通常不需要反向传播。
# view(-1, nc) :将 pred_scores 的形状从 (batch_size, num_queries, num_classes) 变为 (batch_size * num_queries, num_classes) 。 -1 表示自动计算该维度的大小。
pred_scores = pred_scores.detach().view(-1, nc)
# 根据 use_fl 属性的值,选择使用sigmoid函数或softmax函数对预测得分进行处理。
# F.sigmoid(pred_scores) :如果 use_fl 为 True ,使用sigmoid函数对预测得分进行处理。sigmoid函数将得分映射到 (0, 1) 区间,适用于Focal Loss。
# F.softmax(pred_scores, dim=-1) :如果 use_fl 为 False ,使用softmax函数对预测得分进行处理。softmax函数将得分归一化为概率分布,适用于多分类问题。
pred_scores = F.sigmoid(pred_scores) if self.use_fl else F.softmax(pred_scores, dim=-1)
# [batch_size * num_queries, 4]
# 将 预测的边界框 pred_bboxes 从形状 (batch_size, num_queries, 4) 展平为形状 (batch_size * num_queries, 4) 。
# detach() :调用 detach() 方法,使 pred_bboxes 不再参与梯度计算。
# view(-1, 4) :将 pred_bboxes 的形状从 (batch_size, num_queries, 4) 变为 (batch_size * num_queries, 4) 。 -1 表示自动计算该维度的大小。
pred_bboxes = pred_bboxes.detach().view(-1, 4)
# 这三行代码的主要作用是将预测的得分和边界框进行展平处理,以便在批量中计算成本矩阵。展平预测得分:将预测得分从形状 (batch_size, num_queries, num_classes) 展平为形状 (batch_size * num_queries, num_classes) 。处理预测得分:根据是否使用Focal Loss,选择使用sigmoid函数或softmax函数对预测得分进行处理。展平预测边界框:将预测边界框从形状 (batch_size, num_queries, 4) 展平为形状 (batch_size * num_queries, 4) 。这些预处理步骤为后续的成本计算提供了方便,使得可以在批量中高效地计算分类成本、L1距离成本和GIoU成本。
# 这段代码继续 HungarianMatcher 类的 forward 方法,用于计算分类成本、L1距离成本和GIoU成本。
# Compute the classification cost
# 从 预测得分 中选取与 真实类别 gt_cls 对应的得分。
# pred_scores[:, gt_cls] : pred_scores 的形状为 (batch_size * num_queries, num_classes) , gt_cls 的形状为 (num_gt,) 。通过索引操作,选取每个预测得分中与真实类别对应的得分,结果形状为 (batch_size * num_queries, num_gt) 。
pred_scores = pred_scores[:, gt_cls]
# 根据是否使用Focal Loss,计算分类成本。
# 如果为 True ,使用Focal Loss计算分类成本。
if self.use_fl:
# 计算负样本的成本。
neg_cost_class = (1 - self.alpha) * (pred_scores**self.gamma) * (-(1 - pred_scores + 1e-8).log())
# 计算正样本的成本。
pos_cost_class = self.alpha * ((1 - pred_scores) ** self.gamma) * (-(pred_scores + 1e-8).log())
# 正样本成本减去负样本成本,得到 最终的分类成本 。
cost_class = pos_cost_class - neg_cost_class
# 如果不使用Focal Loss,直接取 预测得分的负值 作为 分类成本 。
else:
cost_class = -pred_scores
# Compute the L1 cost between boxes
# 计算预测边界框和真实边界框之间的 L1距离成本 。
# pred_bboxes.unsqueeze(1) :将预测边界框的形状从 (batch_size * num_queries, 4) 变为 (batch_size * num_queries, 1, 4) 。
# gt_bboxes.unsqueeze(0) :将真实边界框的形状从 (num_gt, 4) 变为 (1, num_gt, 4) 。
# abs().sum(-1) :计算两个张量的绝对差值,并在最后一个维度上求和,结果形状为 (batch_size * num_queries, num_gt) 。
cost_bbox = (pred_bboxes.unsqueeze(1) - gt_bboxes.unsqueeze(0)).abs().sum(-1) # (bs*num_queries, num_gt)
# Compute the GIoU cost between boxes, (bs*num_queries, num_gt)
# 计算预测边界框和真实边界框之间的 GIoU成本 。
# bbox_iou :调用 bbox_iou 函数计算边界框的IoU(交并比),并设置 xywh=True 表示边界框的格式为 (x, y, w, h) , GIoU=True 表示计算GIoU。
# squeeze(-1) :去除最后一个维度的冗余,结果形状为 (batch_size * num_queries, num_gt) 。
# 1.0 - ... :GIoU的值范围为 [-1, 1] ,1减去GIoU值,使得成本值范围为 [0, 2] ,成本越小表示匹配度越高。
cost_giou = 1.0 - bbox_iou(pred_bboxes.unsqueeze(1), gt_bboxes.unsqueeze(0), xywh=True, GIoU=True).squeeze(-1)
# 这段代码的主要作用是计算预测目标和真实目标之间的分类成本、L1距离成本和GIoU成本。分类成本:从预测得分中选取与真实类别对应的得分。根据是否使用Focal Loss,计算分类成本。L1距离成本:计算预测边界框和真实边界框之间的L1距离成本。GIoU成本:计算预测边界框和真实边界框之间的GIoU成本。这些成本将被用于构建最终的成本矩阵,通过匈牙利算法找到最优匹配。这些成本的计算确保了匹配过程不仅考虑了分类的准确性,还考虑了边界框的几何匹配度。
# 这段代码继续 HungarianMatcher 类的 forward 方法,用于构建最终的成本矩阵,并在需要时加入掩码成本。
# Final cost matrix
# 构建最终的 成本矩阵 C ,通过将分类成本、L1距离成本和GIoU成本加权求和得到。
# C :最终的成本矩阵,形状为 (batch_size * num_queries, num_gt) 。
C = (
# 分类成本 乘以对应的权重。
self.cost_gain["class"] * cost_class
# L1距离成本 乘以对应的权重。
+ self.cost_gain["bbox"] * cost_bbox
# GIoU成本 乘以对应的权重。
+ self.cost_gain["giou"] * cost_giou
)
# Compute the mask cost and dice cost
# 如果 with_mask 属性为 True ,则计算 掩码成本 ,并将其加到最终的成本矩阵 C 中。
# self.with_mask :一个布尔值,表示是否考虑掩码信息。
# self._cost_mask(bs, gt_groups, masks, gt_mask) :调用 _cost_mask 方法计算掩码成本。该方法它返回一个形状为 (batch_size * num_queries, num_gt) 的成本矩阵。
# C += ... :将掩码成本加到最终的成本矩阵 C 中。
if self.with_mask:
C += self._cost_mask(bs, gt_groups, masks, gt_mask)
# 这段代码的主要作用是构建最终的成本矩阵 C ,并根据是否考虑掩码信息,加入掩码成本。构建最终成本矩阵:将分类成本、L1距离成本和GIoU成本加权求和,得到最终的成本矩阵 C 。加入掩码成本:如果 with_mask 为 True ,调用 _cost_mask 方法计算掩码成本,并将其加到最终的成本矩阵 C 中。最终的成本矩阵 C 将用于后续的匈牙利算法匹配,找到预测目标和真实目标之间的最优匹配。通过这种方式, HungarianMatcher 类能够综合考虑分类、边界框和掩码信息,实现更准确的目标匹配。
# 这段代码是 HungarianMatcher 类的 forward 方法的最后部分,主要目的是处理成本矩阵中的无效值,执行匈牙利算法进行匹配,并返回匹配结果。
# Set invalid values (NaNs and infinities) to 0 (fixes ValueError: matrix contains invalid numeric entries) 将无效值(NaN 和无穷大)设置为 0(修复 ValueError:矩阵包含无效的数字条目)。
# 将成本矩阵 C 中的NaN值和无穷大值替换为0,避免后续计算出错。
# C.isnan() :生成一个布尔张量,表示 C 中的每个元素是否为NaN。
# C.isinf() :生成一个布尔张量,表示 C 中的每个元素是否为无穷大。
# C.isnan() | C.isinf() :逻辑或操作,生成一个布尔张量,表示 C 中的每个元素是否为NaN或无穷大。
# C[...] = 0.0 :将这些无效值替换为0。
C[C.isnan() | C.isinf()] = 0.0
# 将成本矩阵 C 的形状从 (batch_size * num_queries, num_gt) 变为 (batch_size, num_queries, num_gt) ,并移至CPU端。
# C.view(bs, nq, -1) :将 C 的形状从 (batch_size * num_queries, num_gt) 变为 (batch_size, num_queries, num_gt) 。 -1 表示自动计算该维度的大小。
# .cpu() :将张量移至CPU端,因为 linear_sum_assignment 函数在CPU上运行。
C = C.view(bs, nq, -1).cpu()
# scipy.optimize.linear_sum_assignment(cost_matrix, maximize=False)
# scipy.optimize.linear_sum_assignment() 是 SciPy 库中的一个函数,用于解决线性和分配问题,也称为二分图中的最小权重匹配问题。这个函数可以找到一个最优的分配方案,使得总成本最小化(或最大化)。
# 参数 :
# cost_matrix :一个二维数组(成本矩阵),表示二分图中各节点匹配的代价。
# maximize :一个布尔值,默认为 False 。如果设置为 True ,则计算最大权重匹配;如果为 False ,则计算最小权重匹配。
# 返回值 :
# row_ind :一个数组,包含最优分配方案中行的索引。
# col_ind :一个数组,包含最优分配方案中列的索引。
# 这两个数组一起给出了成本矩阵中每一行与每一列的最优匹配。分配的成本可以通过 cost_matrix[row_ind, col_ind].sum() 计算得出。如果成本矩阵是方阵,返回的行索引将被排序,并且等于 numpy.arange(cost_matrix.shape[0]) 。
# 说明 :
# 线性和分配问题在二分图中也称为最小权重匹配。问题实例由矩阵 C 描述,其中每个 C[i,j] 是匹配第一个部分集合的顶点 i 和第二个集合的顶点 j 的成本。目标是找到一个完整的分配,使得总成本最小化。
# 形式上,让 X 是一个布尔矩阵,其中 X[i,j] = 1 表示行 i 分配给列 j 。那么最优分配的成本为 :
# min∑i∑jCijXij
# 其中,在矩阵 X 为方阵的情况下,每一行恰好分配给一列,每一列恰好分配给一行。
# 该函数还可以解决成本矩阵为矩形的经典分配问题的推广。如果它有比列数多的行数,则不需要将每一行都分配给一列,反之亦然。
# 这个实现是修改的 Jonker-Volgenant 算法,没有初始化。
# 使用 linear_sum_assignment 函数对每个批次的成本矩阵进行匈牙利算法匹配,得到匹配的索引。
# C.split(gt_groups, -1) :将成本矩阵 C 按最后一个维度( num_gt )根据 gt_groups 进行分割,返回一个列表,每个元素是一个张量,表示 每个批次的成本矩阵 。
# linear_sum_assignment(c[i]) :对每个批次的成本矩阵 c[i] 调用 linear_sum_assignment 函数,返回匹配的索引。 linear_sum_assignment 函数返回两个数组,分别表示 预测目标 和 真实目标的匹配索引 。
indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(gt_groups, -1))]
# torch.cumsum(input, dim, *, dtype=None, out=None) → Tensor
# torch.cumsum() 是 PyTorch 中的一个函数,用于计算输入张量(tensor)在指定维度上的累积和。
# 参数 :
# input :输入的张量,可以是任意形状的张量,数据类型通常为数值类型(如 float、int 等)。
# dim :指定沿着哪个维度计算累积和。例如,对于一个二维张量, dim=0 表示沿着行方向(即对每一列进行累积和计算), dim=1 表示沿着列方向(即对每一行进行累积和计算)。
# dtype :可选参数,用于指定输出张量的数据类型。如果不指定,则输出张量的数据类型与输入张量相同。
# out :可选参数,用于指定输出张量。如果指定了 out ,则将结果存储在 out 中,而不是创建一个新的张量。
# 返回值 :
# 返回一个新的张量,其形状与输入张量相同,但每个元素是输入张量在指定维度上从起始位置到当前位置的累积和。
# 计算 真实分组的累积和 ,用于后续调整匹配索引。
# [0, *gt_groups[:-1]] :将 gt_groups 的前 n-1 个元素前面加上0,形成一个新的列表。
# .cumsum_(0) :计算累积和,结果存储在原张量中。这个累积和用于调整匹配索引,确保每个批次的索引是连续的。
gt_groups = torch.as_tensor([0, *gt_groups[:-1]]).cumsum_(0) # (idx for queries, idx for gt)
# 根据 匹配索引 和 真实分组的累积和 ,生成最终的匹配结果列表并返回。
# torch.tensor(i, dtype=torch.long) :将匹配索引 i 转换为 torch.long 类型的张量。
# torch.tensor(j, dtype=torch.long) + gt_groups[k] :将匹配索引 j 转换为 torch.long 类型的张量,并加上累积和 gt_groups[k] ,确保每个批次的索引是连续的。
# for k, (i, j) in enumerate(indices) :遍历每个批次的匹配索引,生成最终的匹配结果列表。
return [
(torch.tensor(i, dtype=torch.long), torch.tensor(j, dtype=torch.long) + gt_groups[k])
for k, (i, j) in enumerate(indices)
]
# 这段代码的主要作用是处理成本矩阵中的无效值,执行匈牙利算法进行匹配,并返回匹配结果。处理无效值:将成本矩阵 C 中的NaN值和无穷大值替换为0。调整成本矩阵形状:将成本矩阵 C 的形状从 (batch_size * num_queries, num_gt) 变为 (batch_size, num_queries, num_gt) ,并移至CPU端。执行匈牙利算法:使用 linear_sum_assignment 函数对每个批次的成本矩阵进行匹配,得到匹配的索引。调整匹配索引:计算真实分组的累积和,调整匹配索引,确保每个批次的索引是连续的。返回匹配结果:生成最终的匹配结果列表并返回。通过这些步骤, HungarianMatcher 类能够高效地将预测目标与真实目标进行匹配,为后续的损失计算等操作提供依据。
# 这段代码实现了 HungarianMatcher 类的前向传播方法 forward 。初始化和预处理:获取输入的形状,处理特殊情况(如没有真实目标)。成本计算:分类成本:根据是否使用Focal Loss计算分类成本。L1距离成本:计算预测边界框和真实边界框之间的L1距离成本。GIoU成本:计算预测边界框和真实边界框之间的GIoU成本。掩码成本:如果考虑掩码信息,计算掩码成本。成本矩阵构建:将不同成本项加权求和,得到最终的成本矩阵。匈牙利算法匹配:使用 linear_sum_assignment 函数对成本矩阵进行匈牙利算法匹配,得到最优匹配索引。结果调整和返回:根据真实分组的累积和调整匹配索引,生成最终的匹配结果列表并返回。通过这些步骤, HungarianMatcher 类能够高效地将预测目标与真实目标进行匹配,为后续的损失计算等操作提供依据。
# 这段代码实现了一个基于匈牙利算法的目标匹配器 HungarianMatcher ,用于在目标检测等任务中将预测的目标与真实目标进行匹配。它考虑了分类成本、边界框的L1距离成本和GIoU成本,并可选地考虑掩码信息。通过匈牙利算法,能够高效地找到最优的匹配结果,为后续的损失计算等操作提供依据。
# This function is for future RT-DETR Segment models 此功能适用于未来的 RT-DETR 细分模型。
# def _cost_mask(self, bs, num_gts, masks=None, gt_mask=None):
# assert masks is not None and gt_mask is not None, 'Make sure the input has `mask` and `gt_mask`'
# # all masks share the same set of points for efficient matching
# sample_points = torch.rand([bs, 1, self.num_sample_points, 2])
# sample_points = 2.0 * sample_points - 1.0
#
# out_mask = F.grid_sample(masks.detach(), sample_points, align_corners=False).squeeze(-2)
# out_mask = out_mask.flatten(0, 1)
#
# tgt_mask = torch.cat(gt_mask).unsqueeze(1)
# sample_points = torch.cat([a.repeat(b, 1, 1, 1) for a, b in zip(sample_points, num_gts) if b > 0])
# tgt_mask = F.grid_sample(tgt_mask, sample_points, align_corners=False).squeeze([1, 2])
#
# with torch.cuda.amp.autocast(False):
# # binary cross entropy cost
# pos_cost_mask = F.binary_cross_entropy_with_logits(out_mask, torch.ones_like(out_mask), reduction='none')
# neg_cost_mask = F.binary_cross_entropy_with_logits(out_mask, torch.zeros_like(out_mask), reduction='none')
# cost_mask = torch.matmul(pos_cost_mask, tgt_mask.T) + torch.matmul(neg_cost_mask, 1 - tgt_mask.T)
# cost_mask /= self.num_sample_points
#
# # dice cost
# out_mask = F.sigmoid(out_mask)
# numerator = 2 * torch.matmul(out_mask, tgt_mask.T)
# denominator = out_mask.sum(-1, keepdim=True) + tgt_mask.sum(-1).unsqueeze(0)
# cost_dice = 1 - (numerator + 1) / (denominator + 1)
#
# C = self.cost_gain['mask'] * cost_mask + self.cost_gain['dice'] * cost_dice
# return C
3.def get_cdn_group(batch, num_classes, num_queries, class_embed, num_dn=100, cls_noise_ratio=0.5, box_noise_scale=1.0, training=False):
# 这段代码定义了一个名为 get_cdn_group 的函数,它用于在目标检测模型的训练过程中生成去噪(denoising)数据。这个函数的主要作用是为模型提供带有噪声的标注数据,以增强模型对噪声的鲁棒性。
# 定义了一个名为 get_cdn_group 的函数,接受八个参数。
# 1.batch :这是一个包含批量数据的字典或对象,其中包含了用于训练模型的必要信息,如标注的类别、边界框等。
# 2.num_classes :一个整数,表示数据集中目标类别的总数。
# 3.num_queries :一个整数,表示模型在解码器中使用的查询(queries)数量,即模型预测目标的数量。
# 4.class_embed :一个张量,包含了类别的嵌入表示,用于将类别标签转换为嵌入空间中的向量。
# 5.num_dn :一个整数,表示去噪样本的数量,默认值为100。去噪样本是在训练过程中添加噪声的样本,用于提高模型的鲁棒性。
# 6.cls_noise_ratio :一个浮点数,表示类别标签噪声的比例,默认值为0.5。这意味着有一半的概率类别标签会被随机替换为其他类别。
# 7.box_noise_scale :一个浮点数,表示边界框噪声的缩放比例,默认值为1.0。这个值用于控制添加到边界框的噪声的大小。
# 8.training :一个布尔值,表示模型是否处于训练模式。如果为 False ,则不生成去噪数据。
def get_cdn_group(
batch, num_classes, num_queries, class_embed, num_dn=100, cls_noise_ratio=0.5, box_noise_scale=1.0, training=False
):
# 获取对比去噪训练组。此函数使用来自地面实况 (gt) 的正样本和负样本创建一个对比去噪训练组。它将噪声应用于类标签和边界框坐标,并返回修改后的标签、边界框、注意掩码和元信息。
"""
Get contrastive denoising training group. This function creates a contrastive denoising training group with positive
and negative samples from the ground truths (gt). It applies noise to the class labels and bounding box coordinates,
and returns the modified labels, bounding boxes, attention mask and meta information.
Args:
batch (dict): A dict that includes 'gt_cls' (torch.Tensor with shape [num_gts, ]), 'gt_bboxes'
(torch.Tensor with shape [num_gts, 4]), 'gt_groups' (List(int)) which is a list of batch size length
indicating the number of gts of each image.
num_classes (int): Number of classes.
num_queries (int): Number of queries.
class_embed (torch.Tensor): Embedding weights to map class labels to embedding space.
num_dn (int, optional): Number of denoising. Defaults to 100.
cls_noise_ratio (float, optional): Noise ratio for class labels. Defaults to 0.5.
box_noise_scale (float, optional): Noise scale for bounding box coordinates. Defaults to 1.0.
training (bool, optional): If it's in training mode. Defaults to False.
Returns:
(Tuple[Optional[Tensor], Optional[Tensor], Optional[Tensor], Optional[Dict]]): The modified class embeddings,
bounding boxes, attention mask and meta information for denoising. If not in training mode or 'num_dn'
is less than or equal to 0, the function returns None for all elements in the tuple.
"""
# 这段代码是 get_cdn_group 函数的一部分,它负责在函数开始时进行一些基本的检查和准备工作。
# 检查两个条件。如果模型不在训练模式( training 为 False )或者去噪样本的数量( num_dn )小于或等于0。如果任一条件为真,则函数返回四个 None 值,表示不进行后续处理。
if (not training) or num_dn <= 0:
return None, None, None, None
# 从输入的 batch 数据中提取 gt_groups ,它包含了每个样本的标注组信息。 gt_groups 通常是一个列表,其中的每个元素代表一个样本中的 标注数量 。
gt_groups = batch["gt_groups"]
# 计算所有样本中 标注的总数 。这是通过对 gt_groups 中的元素求和得到的。
total_num = sum(gt_groups)
# 找出一个批次中 标注数量的最大值 。这是通过对 gt_groups 应用 max 函数得到的。
max_nums = max(gt_groups)
# 如果最大标注数量为0,即没有任何标注数据,函数返回四个 None 值。
if max_nums == 0:
return None, None, None, None
# 计算 每个标注组需要生成的去噪样本组数 。这是通过将去噪样本总数( num_dn )除以最大标注数量( max_nums )得到的。
num_group = num_dn // max_nums
# 如果计算出的 num_group 为0,则将其设置为1。这是为了防止后续操作中出现除以0的错误。
num_group = 1 if num_group == 0 else num_group
# Pad gt to max_num of a batch
# 获取 批次大小 ( bs ),即 样本的数量 。这是通过对 gt_groups 列表的长度进行计数得到的。
bs = len(gt_groups)
# 从 batch 数据中提取标注的类别信息( gt_cls ),它是一个一维张量,包含了 每个标注的类别标签 。
gt_cls = batch["cls"] # (bs*num, )
# 从 batch 数据中提取标注的边界框信息( gt_bbox ),它是一个二维张量,形状为 (bs*num, 4) ,包含了 每个标注的边界框坐标 。
gt_bbox = batch["bboxes"] # bs*num, 4
# 从 batch 数据中提取批次索引( b_idx ),它是一个一维张量,包含了 每个标注所属样本的索引 。
b_idx = batch["batch_idx"]
# 这段代码的主要作用是验证输入条件,准备批次数据,并为后续的去噪数据处理做好准备。如果输入条件不满足或者没有标注数据,函数将提前返回,不执行后续的去噪操作。
# 这段代码是 get_cdn_group 函数中用于生成去噪(denoising)查询的部分,其中包括正查询(positive queries)和负查询(negative queries)。
# Each group has positive and negative queries. 每个组都有正和负的问题。
# gt_cls 是一个一维张量,包含了每个标注的 类别标签 。
# repeat(2 * num_group) 将 gt_cls 中的每个元素重复 2 * num_group 次。这样做是为了为每个标注生成两个查询:一个正查询和一个负查询。
# 结果是一个形状为 (2 * num_group * bs * num, ) 的一维张量,其中 bs 是批次大小, num 是每个样本中标注的数量。
dn_cls = gt_cls.repeat(2 * num_group) # (2*num_group*bs*num, )
# gt_bbox 是一个二维张量,形状为 (bs * num, 4) ,包含了每个标注的 边界框坐标 。
# repeat(2 * num_group, 1) 将 gt_bbox 中的每一行(每个标注的边界框)重复 2 * num_group 次,每一列(边界框的四个坐标)保持不变。
# 结果是一个形状为 (2 * num_group * bs * num, 4) 的二维张量。
dn_bbox = gt_bbox.repeat(2 * num_group, 1) # 2*num_group*bs*num, 4
# b_idx 是一个一维张量,包含了每个标注 所属样本的索引 。
# repeat(2 * num_group) 将 b_idx 中的每个元素重复 2 * num_group 次。
# view(-1) 将重复后的张量展平成一维张量。
# 结果是一个形状为 (2 * num_group * bs * num, ) 的一维张量。
dn_b_idx = b_idx.repeat(2 * num_group).view(-1) # (2*num_group*bs*num, )
# 这段代码的目的是为每个标注生成两组查询:正查询和负查询。正查询用于指导模型学习如何从特征中检测目标,而负查询则用于增强模型的判别能力,使其能够区分目标和非目标。通过重复标注数据,可以在不实际增加数据集大小的情况下增加模型训练的样本数量,这是一种数据增强技术。
# 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 。
# 这段代码继续处理 get_cdn_group 函数中的去噪查询,包括生成负样本掩码、添加类别噪声和边界框噪声。
# Positive and negative mask
# (bs*num*num_group, ), the second total_num*num_group part as negative samples
# 生成负样本掩码。创建一个从 total_num * num_group 开始的整数序列,长度为 total_num * num_group ,表示负样本的索引。这个序列被用来区分正样本和负样本。
# 这行代码在 get_cdn_group 函数中用于生成负样本的索引( neg_idx )。
# torch.arange(total_num * num_group, dtype=torch.long, device=gt_bbox.device) :
# torch.arange 函数生成一个从0开始到 total_num * num_group 结束(不包括结束值)的整数序列。
# dtype=torch.long 指定生成的序列的数据类型为长整型。
# device=gt_bbox.device 指定生成的序列位于与 gt_bbox 张量相同的设备上,这可以是 CPU 或 GPU。
# + num_group * total_num :将上述生成的序列中的每个元素加上 num_group * total_num ,这样可以得到负样本的索引。 这个操作实际上是将正样本的索引范围(从0到 total_num * num_group - 1 )转换为负样本的索引范围。
# neg_idx :变量 neg_idx 存储了负样本的索引,这些索引将用于后续的噪声添加过程,特别是在为类别和边界框添加噪声时。
# 这行代码的作用是确定哪些查询应该被视为负样本(即不对应任何真实目标的查询),这是在去噪训练中模拟噪声数据的一部分。负样本的索引将用于在后续步骤中对这些查询的类别标签和边界框进行随机替换,以此来训练模型识别和忽略噪声的能力。
neg_idx = torch.arange(total_num * num_group, dtype=torch.long, device=gt_bbox.device) + num_group * total_num
# 添加类别噪声。
# 如果类别噪声比例大于0,则执行以下操作。
if cls_noise_ratio > 0:
# Half of bbox prob
# 生成一个与 dn_cls 形状相同的随机张量,每个元素小于 cls_noise_ratio * 0.5 的概率,用于确定哪些类别标签需要被替换。
mask = torch.rand(dn_cls.shape) < (cls_noise_ratio * 0.5)
# 获取需要替换的类别标签的索引。
idx = torch.nonzero(mask).squeeze(-1)
# Randomly put a new one here
# torch.randint_like(input, low=0, high, *, dtype=None, layout=None, device=None, requires_grad=False, memory_format=torch.preserve_format) → Tensor
# torch.randint_like 是 PyTorch 中的一个函数,它用于生成一个与给定输入张量形状相同,并且填充了指定范围内的均匀分布随机整数的张量。
# 参数说明 :
# input :一个已有的张量,其形状将被用来生成输出张量的形状。
# low :整数随机数的下限(包含),默认为0。
# high :整数随机数的上限(不包含)。
# dtype :输出张量的数据类型。如果为 None ,则使用输入张量的数据类型。
# layout :输出张量的内存布局。如果为 None ,则使用输入张量的内存布局。
# device :输出张量所在的设备。如果为 None ,则使用输入张量的设备。
# requires_grad :是否需要计算梯度,通常设置为 False 。
# memory_format :内存格式,通常设置为 torch.preserve_format 。
# 返回值 :
# 一个与输入张量 input 形状相同,并且填充了在 [low, high) 范围内均匀分布的随机整数的张量。
# 这个函数常用于在模型训练中添加随机性,例如在数据增强或者随机采样等场景。通过 torch.randint_like ,可以确保生成的随机整数张量与某个特定张量具有相同的形状,这在处理特定形状的数据时非常有用。
# 为这些索引生成新的随机类别标签。
new_label = torch.randint_like(idx, 0, num_classes, dtype=dn_cls.dtype, device=dn_cls.device)
# 将原始类别标签替换为新的随机类别标签。
dn_cls[idx] = new_label
# 添加边界框噪声。
# 如果边界框噪声比例大于0,则执行以下操作。
if box_noise_scale > 0:
# 将边界框从中心点宽度高度(x, y, w, h)格式转换为左上角和右下角坐标(x1, y1, x2, y2)格式。
known_bbox = xywh2xyxy(dn_bbox)
# 计算边界框的宽度和高度的一半,然后扩展到四个坐标,最后乘以噪声比例,得到噪声差异。
# 这行代码是在处理边界框噪声时使用的,它计算了边界框的宽度和高度的一半,然后扩展到四个坐标,最后乘以噪声比例。
# dn_bbox[..., 2:] :从 dn_bbox 张量中选取最后一个维度的最后两个元素,即每个边界框的宽度和高度(假设 dn_bbox 的形状为 [N, 4] ,其中 N 是边界框的数量,4 代表 x, y, w, h )。
# * 0.5 :将宽度和高度乘以 0.5,得到边界框的一半宽度和高度。
# repeat(1, 2) :将得到的宽度和高度沿着第二个维度(维度索引为 1)重复两次,这样做是为了将宽度和高度分别应用到 x 和 y 坐标上,扩展成四个坐标。
# * box_noise_scale :将上一步得到的结果乘以 box_noise_scale ,这是一个控制噪声大小的缩放比例。
# diff 张量包含了每个边界框宽度和高度的一半,并且这些值被扩展和缩放到四个坐标上,用于后续添加噪声。这个计算结果将被用来生成随机噪声,这些噪声将被添加到边界框的坐标上,以模拟训练数据中的噪声,增强模型对噪声的鲁棒性。
diff = (dn_bbox[..., 2:] * 0.5).repeat(1, 2) * box_noise_scale # 2*num_group*bs*num, 4
# 生成一个与 dn_bbox 形状相同的随机张量,元素值为1或-1,用于控制噪声的方向。
rand_sign = torch.randint_like(dn_bbox, 0, 2) * 2.0 - 1.0
# 生成一个与 dn_bbox 形状相同的随机张量,用于生成噪声。
rand_part = torch.rand_like(dn_bbox)
# 确保负样本的噪声部分大于1,这样在应用 logistic 函数时,这些值会被映射到接近0或1的区域,模拟负样本。
rand_part[neg_idx] += 1.0
# 将噪声的方向应用到噪声大小上。
rand_part *= rand_sign
# 将噪声应用到边界框上。
known_bbox += rand_part * diff
# 将边界框的值限制在0到1之间。
known_bbox.clip_(min=0.0, max=1.0)
# 将边界框从左上角和右下角坐标格式转换回中心点宽度高度格式。
dn_bbox = xyxy2xywh(known_bbox)
# torch.logit(input, eps=1e-6) → Tensor
# torch.logit 是 PyTorch 中的一个函数,它用于计算输入张量的 logit 函数值,即逆 sigmoid 函数。logit 函数定义为: logit(x) = log(x/(1-x)) 。
# 参数说明 :
# input :输入张量,其元素值应该在区间 (0, 1) 内,因为 logit 函数在0和1处未定义。
# eps :一个很小的正数,用于避免除以零的错误。它确保了 1 - x 不会等于零。
# 返回值 :
# 一个与输入张量 input 形状相同,并且填充了输入张量每个元素的 logit 函数值的张量。
# 这个函数常用于将概率值转换为 logit 空间,这在某些机器学习模型中很有用,例如在二分类问题中,logit 函数可以将概率值映射到实数线上,从而可以用于计算损失函数。
# 对边界框应用 logistic 函数,使其值在0和1之间, eps 参数用于防止数值不稳定。
dn_bbox = torch.logit(dn_bbox, eps=1e-6) # inverse sigmoid
# 这段代码通过添加类别噪声和边界框噪声,为模型提供了更加鲁棒的训练数据,帮助模型在面对不完美或噪声数据时更好地进行目标检测。
# 这段代码是 get_cdn_group 函数中的一部分,用于准备去噪查询的类别嵌入和边界框,并创建相应的填充张量。
# 计算 总的去噪查询数量 ,这是基于每个样本中标注数量的最大值( max_nums )、每组的查询数量( num_group )以及正负样本的总数(乘以2)。
# 这行代码计算了去噪查询的总数( num_dn ),它是基于每个样本中标注数量的最大值( max_nums )、每组的查询数量( num_group )以及正负样本的总数(乘以2)。
# max_nums :这是批次中每个样本的标注数量的最大值。例如,如果一个批次中有3个样本,分别有2、3和4个标注,那么 max_nums 就是4。
# num_group :这是每个标注组需要生成的去噪样本组数。例如,如果 num_dn 是100, max_nums 是4,那么 num_group 就是 100 // 4 = 25 。如果 num_group 计算结果为0,则会被设置为1。
# 2 * num_group :由于每个标注组需要生成正样本和负样本,所以乘以2。例如,如果 num_group 是25,那么每个标注组将生成50个去噪查询(25个正样本和25个负样本)。
# int(max_nums * 2 * num_group) :最后,将上述结果转换为整数,得到去噪查询的总数。例如,如果 max_nums 是4, num_group 是25,那么 num_dn 就是 int(4 * 2 * 25) = 200 。
# 这个计算结果 num_dn 将用于后续的去噪数据准备,包括生成去噪查询的类别嵌入和边界框,以及创建相应的填充张量。通过这种方式,模型可以学习在存在噪声的情况下如何准确地检测目标。
num_dn = int(max_nums * 2 * num_group) # total denoising queries
# class_embed = torch.cat([class_embed, torch.zeros([1, class_embed.shape[-1]], device=class_embed.device)])
# 使用 class_embed (类别嵌入矩阵)和 dn_cls (去噪类别标签)来获取 去噪查询的类别嵌入 。结果是一个形状为 (batch_size * num_annotations * 2 * num_group, embedding_dim) 的张量。
# 这行代码在 get_cdn_group 函数中用于获取去噪查询的类别嵌入。
# class_embed :这是一个张量,包含了所有类别的嵌入表示。它的形状通常是 [num_classes, embedding_dim] ,其中 num_classes 是类别的总数, embedding_dim 是嵌入的维度。
# dn_cls :这是一个一维张量,包含了去噪查询的类别标签。它的形状是 [2 * num_group * bs * num] ,其中 bs 是批次大小, num 是每个样本的标注数量, num_group 是每个标注组的去噪样本组数。
# dn_cls_embed = class_embed[dn_cls] :这行代码使用 dn_cls 中的类别标签作为索引,从 class_embed 中获取对应的类别嵌入。结果是一个形状为 [2 * num_group * bs * num, embedding_dim] 的张量,包含了所有去噪查询的类别嵌入。
# 这个操作的目的是将去噪查询的类别标签转换为嵌入空间中的向量,这些向量将被用于后续的模型训练过程中,以提供模型所需的类别信息。通过使用类别嵌入,模型可以更好地理解和处理类别之间的关系,从而提高目标检测的准确性。
dn_cls_embed = class_embed[dn_cls] # bs*num * 2 * num_group, 256
# 创建一个形状为 (batch_size, total_denoising_queries, embedding_dim) 的零填充张量,用于存储去噪查询的 类别嵌入 。
padding_cls = torch.zeros(bs, num_dn, dn_cls_embed.shape[-1], device=gt_cls.device)
# 创建一个形状为 (batch_size, total_denoising_queries, 4) 的零填充张量,用于存储去噪查询的 边界框 。
padding_bbox = torch.zeros(bs, num_dn, 4, device=gt_bbox.device)
# 创建一个索引映射,它是一个一维张量,包含了每个样本中 标注的索引 。
map_indices = torch.cat([torch.tensor(range(num), dtype=torch.long) for num in gt_groups])
# 为每个样本生成 正样本的索引 。这是通过将 map_indices 与 max_nums 的倍数相加来实现的,以确保每个样本的正样本索引是唯一的。
pos_idx = torch.stack([map_indices + max_nums * i for i in range(num_group)], dim=0)
# 扩展索引映射,包括正样本和负样本的索引。
map_indices = torch.cat([map_indices + max_nums * i for i in range(2 * num_group)])
# 将去噪查询的 类别 嵌入填充到 padding_cls 张量中。这里使用 dn_b_idx (去噪批次索引)和 map_indices (扩展的索引映射)作为索引。
padding_cls[(dn_b_idx, map_indices)] = dn_cls_embed
# 将去噪查询的 边界框 填充到 padding_bbox 张量中。同样使用 dn_b_idx 和 map_indices 作为索引。
padding_bbox[(dn_b_idx, map_indices)] = dn_bbox
# 这段代码的主要作用是为去噪查询准备类别嵌入和边界框的数据,并创建填充张量以存储这些信息。这些填充张量将被用于后续的模型训练过程中,以提供模型所需的去噪数据。通过这种方式,模型可以学习在存在噪声的情况下如何准确地检测目标。
# 这段代码是 get_cdn_group 函数的最后部分,它负责创建注意力掩码( attn_mask )和去噪元数据( dn_meta ),并返回所有准备好的数据。
# 计算目标大小( tgt_size ),这是去噪查询数量( num_dn )和模型查询数量( num_queries )的总和。
tgt_size = num_dn + num_queries
# 创建一个形状为 [tgt_size, tgt_size] 的布尔类型零张量,用于作为注意力掩码。这个掩码将用于在 Transformer 模型中控制注意力权重,以防止某些查询看到不应该看到的信息。
attn_mask = torch.zeros([tgt_size, tgt_size], dtype=torch.bool)
# Match query cannot see the reconstruct
# 将掩码中对应于模型查询(match query)的行和对应于去噪查询(reconstruct)的列设置为 True 。这意味着模型查询不能看到去噪查询的信息。
attn_mask[num_dn:, :num_dn] = True
# Reconstruct cannot see each other
# 遍历每个去噪组( num_group )。
for i in range(num_group):
# 对于第一个去噪组,将掩码中对应于当前组的去噪查询和下一个组的去噪查询之间的区域设置为 True ,以防止它们相互看到。
if i == 0:
attn_mask[max_nums * 2 * i : max_nums * 2 * (i + 1), max_nums * 2 * (i + 1) : num_dn] = True
# 对于最后一个去噪组,将掩码中对应于当前组的去噪查询和之前所有组的去噪查询之间的区域设置为 True 。
if i == num_group - 1:
attn_mask[max_nums * 2 * i : max_nums * 2 * (i + 1), : max_nums * i * 2] = True
# 对于中间的去噪组,将掩码中对应于当前组的去噪查询和前后组的去噪查询之间的区域设置为 True 。
else:
attn_mask[max_nums * 2 * i : max_nums * 2 * (i + 1), max_nums * 2 * (i + 1) : num_dn] = True
attn_mask[max_nums * 2 * i : max_nums * 2 * (i + 1), : max_nums * 2 * i] = True
# 创建一个字典 dn_meta ,包含 去噪的正样本索引 ( dn_pos_idx )、 去噪组的数量 ( dn_num_group )和 去噪查询与模型查询的分割数量 ( dn_num_split )。
dn_meta = {
"dn_pos_idx": [p.reshape(-1) for p in pos_idx.cpu().split(list(gt_groups), dim=1)],
"dn_num_group": num_group,
"dn_num_split": [num_dn, num_queries],
}
# 返回准备好的数据,包括 类别嵌入的填充张量( padding_cls )、 边界框的填充张量 ( padding_bbox )、 注意力掩码 ( attn_mask )和 去噪元数据 ( dn_meta )。这些数据都被转移到了 class_embed 的设备上,以确保所有的计算都在同一个设备上进行。
return (
padding_cls.to(class_embed.device),
padding_bbox.to(class_embed.device),
attn_mask.to(class_embed.device),
dn_meta,
)
# 这段代码的目的是为模型训练准备好所有必要的去噪数据和掩码,以便模型可以学习在存在噪声的情况下进行目标检测。通过这种方式,模型可以提高其对噪声的鲁棒性,并在实际应用中表现得更加稳定。
# 这个函数的主要作用是为模型提供带有噪声的标注数据,以增强模型对噪声的鲁棒性。通过添加类别噪声和边界框噪声,模型可以学习在面对不完美或噪声数据时如何更好地进行目标检测。