目标检测 TaskAlignedAssigner 原理

TaskAlignedAssigner 原理和代码

原理主要是结合预测的分类分数和边界框与真实标注的信息，找出与真实目标最匹配的锚点，为这些锚点分配对应的目标标签、边界框和分数。

TaskAlignedAssigner 是目标检测中用于对齐分类和定位任务的样本分配器。其核心思想是通过综合分类得分和预测框与真实框的 IoU，动态选择最合适的锚点作为正样本。具体步骤如下：

1. 计算 IoU：预测框与真实框之间的 IoU。

2. 提取分类得分：根据真实框的类别，提取对应类别的分类得分。

3. 任务对齐指标：计算每个锚点的任务对齐指标，公式为：
   

其中，$\alpha$ 和 $\beta$ 是超参数，用于平衡分类和定位的重要性。

4. 中心约束：过滤掉锚点中心不在真实框内的候选。

5. 动态 Top-k 选择：对每个真实框，选择任务对齐指标最高的前 $k$ 个锚点作为正样本。

6. 冲突处理：若一个锚点被多个真实框选中，保留指标最高的分配。

import torch
import torch.nn as nndef pairwise_iou(boxes1, boxes2):"""计算两组框之间的 IoU。Args:boxes1 (Tensor): (N, 4) 格式为 xyxyboxes2 (Tensor): (M, 4) 格式为 xyxyReturns:iou (Tensor): (N, M)"""lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])wh = (rb - lt).clamp(min=0)inter = wh[:, :, 0] * wh[:, :, 1]area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])union = area1[:, None] + area2 - interreturn inter / (union + 1e-9)
class TaskAlignedAssigner(nn.Module):def __init__(self, topk=13, alpha=1.0, beta=6.0):super().__init__()# 初始化超参数self.topk = topk    # 每个真实框（GT）最多选择的正样本锚点数量，默认13self.alpha = alpha  # 分类得分的指数权重，用于平衡分类任务的重要性，默认1.0self.beta = beta    # IoU的指数权重，用于平衡定位任务的重要性，默认[e-mail]_grad()def forward(self, cls_scores, bbox_preds, gt_bboxes, gt_labels):"""输入:cls_scores  : (B, num_anchors, num_classes) 模型输出的分类得分（未归一化）bbox_preds : (B, num_anchors, 4)           模型输出的预测框坐标（xyxy格式）gt_bboxes  : (B, num_gts, 4)               真实框坐标（xyxy格式）gt_labels  : (B, num_gts)                  真实框的类别标签（0~num_classes-1）输出:pos_anchor_indices : (batch_idx, anchor_idx) 正样本锚点的批次索引和锚点索引pos_gt_indices     : (gt_idx,)               对应的真实框索引pos_labels         : (label,)               正样本的类别标签"""# 获取输入张量的维度信息batch_size, num_anchors, _ = cls_scores.shape  # B=批次大小, num_anchors=锚点总数device = cls_scores.device  # 设备信息（CPU/GPU）# 初始化分配结果张量（记录每个锚点分配的GT索引和类别）assigned_gt_inds = torch.zeros((batch_size, num_anchors), dtype=torch.long, device=device)  # Shape: (B, num_anchors)，0表示未分配，非零值表示分配的GT索引+1（避免与0冲突）assigned_labels = torch.zeros((batch_size, num_anchors), dtype=torch.long, device=device)  # Shape: (B, num_anchors)，记录分配的类别标签# 逐样本处理（每个批次独立处理）for b in range(batch_size):# 提取当前样本的预测框、真实框和类别标签bbox_pred = bbox_preds[b]  # (num_anchors,4) 当前批次的所有锚点预测框gt_bbox = gt_bboxes[b]     # (num_gts,4)     当前批次的所有真实框gt_label = gt_labels[b]    # (num_gts)       当前批次的所有真实框类别num_gts = gt_bbox.size(0)  # 当前样本的真实框数量# 若当前样本没有真实框，跳过处理if num_gts == 0:continue#################################################################### Step 1: 计算预测框与真实框的 IoU#################################################################### 输入: bbox_pred (num_anchors,4), gt_bbox (num_gts,4)# 输出: iou (num_anchors, num_gts)，每个元素为锚点与GT的IoUiou = pairwise_iou(bbox_pred, gt_bbox)  #################################################################### Step 2: 提取对应真实框类别的分类得分#################################################################### cls_scores[b]的shape为 (num_anchors, num_classes)# gt_label的shape为 (num_gts)，每个元素是对应GT的类别索引# 通过高级索引提取每个锚点在对应GT类别上的得分# 结果scores的shape为 (num_anchors, num_gts)scores = cls_scores[b][:, gt_label]  #################################################################### Step 3: 计算任务对齐指标（分类得分^α * IoU^β）###################################################################alignment_metrics = scores.pow(self.alpha) * iou.pow(self.beta)# alignment_metrics的shape: (num_anchors, num_gts)#################################################################### Step 4: 中心点约束（过滤锚点中心不在GT内部的候选）#################################################################### 计算锚点中心坐标cx = (bbox_pred[:, 0] + bbox_pred[:, 2]) / 2  # (num_anchors,)cy = (bbox_pred[:, 1] + bbox_pred[:, 3]) / 2  # (num_anchors,)# 判断锚点中心是否在GT框内（利用广播机制）# cx[:, None]的shape: (num_anchors,1)# gt_bbox[None, :, 0]的shape: (1, num_gts)# 比较操作后，in_gt的shape为 (num_anchors, num_gts)in_gt = (cx[:, None] >= gt_bbox[None, :, 0]) & \(cx[:, None] <= gt_bbox[None, :, 2]) & \(cy[:, None] >= gt_bbox[None, :, 1]) & \(cy[:, None] <= gt_bbox[None, :, 3])# 将不在GT内的锚点指标置零alignment_metrics *= in_gt.float()  # (num_anchors, num_gts)#################################################################### Step 5: 动态选择每个GT的Top-k锚点###################################################################candidate_metrics = []       # 保存所有候选锚点的指标值（多个张量）candidate_gt_indices = []    # 保存候选锚点对应的GT索引（平铺列表）candidate_anchor_indices = [] # 保存候选锚点的索引（多个张量）# 遍历每个真实框（GT）for gt_idx in range(num_gts):# 提取当前GT对应的所有锚点指标metrics = alignment_metrics[:, gt_idx]  # (num_anchors,)# 筛选有效锚点（指标>0表示中心在GT内且指标非零）valid = metrics > 0  # (num_anchors,)if not valid.any():  # 若没有有效锚点，跳过该GTcontinue# 确定实际选择的Top-k数量（不超过有效锚点数和预设topk）k = min(self.topk, valid.sum().item())# 选择当前GT的Top-k锚点（指标值和索引）topk_metrics, topk_anchors = metrics.topk(k)# 保存结果candidate_metrics.append(topk_metrics)  # 添加一个形状为(k,)的张量candidate_gt_indices.extend([gt_idx] * k)  # 扩展k个gt_idx元素candidate_anchor_indices.append(topk_anchors)  # 添加形状为(k,)的锚点索引# 若当前样本无候选锚点，跳过后续处理if not candidate_metrics:continue#################################################################### Step 6: 合并候选并排序#################################################################### 合并所有候选指标candidate_metrics = torch.cat(candidate_metrics)  # (total_candidates,)# 转换GT索引为张量（total_candidates,)candidate_gt_indices = torch.tensor(candidate_gt_indices, dtype=torch.long, device=device)# 合并所有候选锚点索引candidate_anchor_indices = torch.cat(candidate_anchor_indices)  # (total_candidates,)# 按指标降序排序（从高到低）sorted_idx = candidate_metrics.argsort(descending=True)  # (total_candidates,)candidate_gt_indices = candidate_gt_indices[sorted_idx]  # 按排序调整GT索引candidate_anchor_indices = candidate_anchor_indices[sorted_idx]  # 按排序调整锚点索引#################################################################### Step 7: 分配正样本（解决冲突，高优先级指标优先）###################################################################assigned_mask = torch.zeros(num_anchors, dtype=torch.bool, device=device)  # 标记锚点是否已被分配# 按排序后的顺序遍历候选锚点for anchor_idx, gt_idx in zip(candidate_anchor_indices, candidate_gt_indices):if not assigned_mask[anchor_idx]:# 记录分配的GT索引（+1避免与0冲突）assigned_gt_inds[b, anchor_idx] = gt_idx + 1  # 记录分配的类别标签assigned_labels[b, anchor_idx] = gt_label[gt_idx]  # 标记该锚点已分配assigned_mask[anchor_idx] = True  #################################################################### Step 8: 提取最终正样本信息#################################################################### 生成正样本的掩码（assigned_gt_inds > 0表示已分配）pos_mask = assigned_gt_inds > 0  # (B, num_anchors)# 获取正样本的批次索引和锚点索引（非零元素的坐标）pos_anchor_indices = pos_mask.nonzero(as_tuple=True)  # 格式为 (batch_indices, anchor_indices)，例如：(tensor([0,0,1]), tensor([5,8,3]))# 获取对应的GT索引（需减去1恢复原始索引）pos_gt_indices = assigned_gt_inds[pos_mask] - 1  # (num_pos_samples,)# 获取正样本的类别标签pos_labels = assigned_labels[pos_mask]  # (num_pos_samples,)return pos_anchor_indices, pos_gt_indices, pos_labels

使用示例

if __name__ == "__main__":# 参数设置batch_size = 2num_anchors = 100   # 锚点数量num_classes = 20     # 类别数num_gts = 3         # 每个样本的 GT 数量# 模拟数据cls_scores = torch.rand(batch_size, num_anchors, num_classes)  # 随机分类得分bbox_preds = torch.rand(batch_size, num_anchors, 4) * 100      # 随机预测框（xyxy）gt_bboxes = torch.rand(batch_size, num_gts, 4) * 100           # 随机真实框（xyxy）gt_labels = torch.randint(0, num_classes, (batch_size, num_gts)) # 随机 GT 类别标签# 初始化分配器assigner = TaskAlignedAssigner(topk=5, alpha=1.0, beta=6.0)# 分配正样本pos_anchors, pos_gts, pos_labels = assigner(cls_scores, bbox_preds, gt_bboxes, gt_labels)# 输出结果print("正样本锚点索引:", pos_anchors)  # 格式为 (batch_idx, anchor_idx) 的元组print("对应真实框索引:", pos_gts)     # 每个正样本锚点对应的 GT 索引print("正样本标签:", pos_labels)      # 每个正样本锚点的类别标签