OpenPCDet系列 | 5.4.3 DenseHead中的AxisAlignedTargetAssigner正负样本分配模块

AxisAlignedTargetAssigner模块

TargetAssigner处理，也就是正负样本的分配问题，算是整个检测算法中一个比较核心的问题，也比较重要。

在AnchorHeadSingle模块总，对分类、回归、方向预测三个特征矩阵预测完之后，随机就调用self.assign_targets函数来对在基类中生成的anchor进行正负样本的匹配，构建出box_cls_labels、box_reg_targets、reg_weights来进行损失函数的构建。而这里的调用self.assign_targets函数其实是调用基类的assign_targets函数，最后再跳转到AxisAlignedTargetAssigner模块的assign_targets函数中。调用关系如下所示：

---

assign_targets处理流程

1. 提取有效gt信息

对于当前一个batch的数据，gt信息是进行填充后再拼接在一起的，所以存在0填充的部分。那么，在进行后续处理的前提是，先对gt的填充信息进行去除。提取每个点云帧的有效gt信息以及有效类别信息，保留非零项。

cur_gt = gt_boxes[k]    # 提取第k个点云帧gt，然后提取非零信息，去除非0无效信息 [44, 7] -> [38, 7]
cnt = cur_gt.__len__() - 1      # 43
while cnt > 0 and cur_gt[cnt].sum() == 0:
    cnt -= 1
cur_gt = cur_gt[:cnt + 1]   # 提取当前第k点云帧的有效gt信息，保留非零项
cur_gt_classes = gt_classes[k][:cnt + 1].int()   # 提取当前第k点云帧有效gt类别

2. 提取需要处理的类别信息

由于当前的gt信息包含了3个类别： [‘Car’, ‘Pedestrian’, ‘Cyclist’]，现在需要对着三个类别进行分别处理。也就是利用掩码矩阵，分别获取每个当前需要处理的类别，同时单独获取当前需要处理的gt信息，然后传入到assign_targets_single函数中进行处理，这个函数丶作用是针对某一个点云帧中的每一个类别anchors和gt信息，计算前景和背景的anchor类别，box编码以及回归的权重。

for anchor_class_name, anchors in zip(self.anchor_class_names, all_anchors):    # 对每个类别及其配置anchor进行依次处理
    if cur_gt_classes.shape[0] > 1:
        mask = torch.from_numpy(self.class_names[cur_gt_classes.cpu() - 1] == anchor_class_name)    # 获取类别为'Car'的掩码矩阵:[True, True, ..., False]
    else:
        mask = torch.tensor([self.class_names[c - 1] == anchor_class_name
                             for c in cur_gt_classes], dtype=torch.bool)

    if self.use_multihead:  # False
        anchors = anchors.permute(3, 4, 0, 1, 2, 5).contiguous().view(-1, anchors.shape[-1])
        selected_classes = cur_gt_classes[mask]
    else:
        feature_map_size = anchors.shape[:3]   # zyx: (1, 248, 216)
        anchors = anchors.view(-1, anchors.shape[-1])   # (107136,7) 107136=1x248x216x1x2
        selected_classes = cur_gt_classes[mask]     # 被选择的类别 (14, )

    single_target = self.assign_targets_single(
        anchors,        # reshape后的anchor矩阵 (107136,7)
        cur_gt[mask],   # 根据当前类别的掩码矩阵选择当前处理的类别gt信息  (38, 7) -> (14, 7)
        gt_classes=selected_classes,    # 当前处理的类别信息 (14, )  [1,1,1, ..., 1,1]]
        matched_threshold=self.matched_thresholds[anchor_class_name],       # 当前处理类别的正样本阈值
        unmatched_threshold=self.unmatched_thresholds[anchor_class_name]    # 当前处理类别的负样本阈值
    )
    target_list.append(single_target)

也就是说，根据掩码矩阵来挑选出对应类别的anchor设置，以及对应类别的gt信息，然后使用assign_targets_single函数进行后续处理。在assign_targets_single函数中完成对当前类别的anchor进行类别赋值以及与其匹配的gt编码信息赋值，还设置了其回归的权重为1。由于这里设置了需要预测3个类别，所以这个函数对应每个点云帧场景会运行3次，依次处理好每个类别。

3. 帧信息整合

对某个点云帧场景的类别信息提取之后，列表信息如下所示：

对其进行合并并拼接起来：

先reshape，再合并，然后再reshape。最后得到如下的结果：

对于每个点云帧都进行信息的整合处理，然后将当前点云场景处理结果分别追加到对应列表中

bbox_targets.append(target_dict['box_reg_targets'])
cls_labels.append(target_dict['box_cls_labels'])
reg_weights.append(target_dict['reg_weights'])

4. 批信息整合

对于整个batch的点云帧信息都处理完之后，其数据结构如下所示：

现在分别对其进行堆叠层一个tensor矩阵处理，随后保存在键值value中：

# 将每个点云帧处理的结果进行stack堆叠
bbox_targets = torch.stack(bbox_targets, dim=0)    # (16, 321408, 7)
cls_labels = torch.stack(cls_labels, dim=0)        # (16, 321408)
reg_weights = torch.stack(reg_weights, dim=0)      # (16, 321408)

all_targets_dict = {
    
    
    'box_cls_labels': cls_labels,
    'box_reg_targets': bbox_targets,
    'reg_weights': reg_weights

}

返回的数据维度如下所示：

自此，完成了对每个点云帧的anchor正样本分配。最后，这个字典的信息会返回到AnchorHeadSingle函数中，保存在self.forward_ret_dict这个字典中，后续就会利用这个字典来进行损失的计算。

在这之后就是进行损失函数的计算：self.get_training_loss()

---

assign_targets_single处理流程

1. 构建每个anchor的正负样本label分配

首先，对于传进来的当前类别的gt信息以及当前类别的生成的anchor，可以进行一个iou3d的计算。也就是先计算anchor和gt之间的iou。

# 1.计算gt和anchors之间的overlap
anchor_by_gt_overlap = iou3d_nms_utils.boxes_iou3d_gpu(anchors[:, 0:7], gt_boxes[:, 0:7]) \
    if self.match_height else box_utils.boxes3d_nearest_bev_iou(anchors[:, 0:7], gt_boxes[:, 0:7])  # 计算anchor和gt之间的iou (107136, 14)

根据这个anchor和gt之间的iou矩阵，可以分别获得与每个anchor最匹配的gt索引以及数值，也可以获得与每个gt最匹配的anchor索引以及数值。其中有可能出现某个gt没有找到与之有任何重叠的anchor，那么最匹配的iou数值为0，此时将其赋值为-1.

# 找到每个anchor最匹配的gt的索引和iou
anchor_to_gt_argmax = anchor_by_gt_overlap.argmax(dim=1)    # (107136，）找到每个anchor最匹配的gt的索引
anchor_to_gt_max = anchor_by_gt_overlap[torch.arange(num_anchors, device=anchors.device), anchor_to_gt_argmax]  # （107136，）找到每个anchor最匹配的gt的iou

# 提取最匹配的anchor，避免没有anchor满足索设定的阈值
# gt_to_anchor_argmax = torch.from_numpy(anchor_by_gt_overlap.cpu().numpy().argmax(axis=0)).cuda()
gt_to_anchor_argmax = anchor_by_gt_overlap.argmax(dim=0)     # (14,) 找到每个gt最匹配anchor的索引
gt_to_anchor_max = anchor_by_gt_overlap[gt_to_anchor_argmax, torch.arange(num_gt, device=anchors.device)]   # （14，）找到每个gt最匹配anchor的iou
empty_gt_mask = gt_to_anchor_max == 0   # 如果最匹配iou为0，表示某个gt没有与之匹配的anchor
gt_to_anchor_max[empty_gt_mask] = -1    # 没有与之匹配的anchor在iou值中设置为-1

接着，根据这一系列最大iou匹配的值，可以找到满足这个最大iou的每个anchor。具体来说，以gt为基础，逐个anchor对应，比如第一个gt的最大iou为0.9，则在所有anchor中找iou为0.9的anchor。对于这些anchor，在labels的对应位置中为其分配类别信息（此类别信息就是当前处理的类别信息），同时也记录为其分配的gt索引。

nchors_with_max_overlap = (anchor_by_gt_overlap == gt_to_anchor_max).nonzero()[:, 0]   # 找到满足最大iou的每个anchor
gt_inds_force = anchor_to_gt_argmax[anchors_with_max_overlap]   # 找到最大iou的gt索引
labels[anchors_with_max_overlap] = gt_classes[gt_inds_force]    # 将gt的类别赋值到对应的anchor的label中 (107136,)
gt_ids[anchors_with_max_overlap] = gt_inds_force.int()          # 将gt的索引赋值到对应的anchor的gt_id中 (107136,)

以上的操作是处理最值以避免没有符合满足阈值设定的iou。接下来还会对满足正样本阈值的anchor进行label和gt的分配。

对于每个anchor与gt的最大iou值，如果超过设定的正样本阈值范围，比如这里的0.6，都会根据这个阈值掩码将符合的anchor挑选出来。然后在labels中对应符合阈值的anchor设置其类别，同时也设置其分配的gt索引。一般情况下，符合最值的anchor的iou匹配只能设置十来个，然后满足阈值的anchor的iou匹配只有百来个。

# 这里应该对labels和gt_ids的操作应该包含了上面的anchors_with_max_overlap
pos_inds = anchor_to_gt_max >= matched_threshold        # 找到最匹配的anchor中iou大于给定阈值的mask (107136,)
gt_inds_over_thresh = anchor_to_gt_argmax[pos_inds]     # 找到最匹配的anchor中iou大于给定阈值的gt的索引 (104,)
labels[pos_inds] = gt_classes[gt_inds_over_thresh]      # 将pos anchor对应gt的类别赋值到对应的anchor的label中 (107136,)
gt_ids[pos_inds] = gt_inds_over_thresh.int()            # 将pos anchor对应gt的索引赋值到对应的anchor的gt_id中 (107136,)

此时，由于已经设置了label的anchor就是正样本，就可以分别找到前景anchor和背景anchor的索引。

bg_inds = (anchor_to_gt_max < unmatched_threshold).nonzero()[:, 0]  # 找到背景anchor索引 (106874，)
fg_inds = (labels > 0).nonzero()[:, 0]      # 找到前景点的索引 (104,)

最后，将labels中的背景anchor类别设置为0.

labels[bg_inds] = 0     # 将背景点的label赋值为0

这时候，就完成了anchor的类别分配操作，一般前景anchor的数量还是比较少的。一个类别中只有百个这样的数量级。相比之下，anchor的数量级是十万。

2. 构建每个anchor的正负样本编码信息bbox_targets分配

对于负样本的anchor预测编码信息设置为0，只对正样本的anchor进行赋值处理。

首先，基于上述的操作已经获取到了最大值iou以及满足设定阈值的正样本anchor的索引，根据真个anchor索引可以获得与其iou最大的gt索引，那么根据gt的索引就可以获取对应的gt信息。同时，根据这个anchor索引也可以获取对应的生成的anchor信息。那么，根据anchor和其对应的gt信息，就可以进行所需预测的编码处理，所获得的编码存储在bbox_targets的前景点索引位置。这里的编码操作是通过self.box_coder.encode_torch实现的。代码如下所示：

# 2. 构建正样本anchor需要预测拟合的编码信息（负样本anchor全部设置为0）
bbox_targets = anchors.new_zeros((num_anchors, self.box_coder.code_size))   # (107136,7)
if len(gt_boxes) > 0 and anchors.shape[0] > 0:
    fg_gt_boxes = gt_boxes[anchor_to_gt_argmax[fg_inds], :]     # 提取前景对应的gt box信息 (104, 7)
    fg_anchors = anchors[fg_inds, :]    # 提取前景anchor (104, 7)
    bbox_targets[fg_inds, :] = self.box_coder.encode_torch(fg_gt_boxes, fg_anchors)    # 编码gt和前景anchor，并赋值到bbox_targets的对应位置

论文中的回归编码方式如下：

这部分的具体代码见ResidualCoder模块，

3. 构建每个anchor的回归权重

这里的回归权重只针对前景anchor，赋值为1.背景的anchor赋值为0.

# 3. 构建正负样本回归权重，其中背景anchor权重为0，前景anchor权重为1
reg_weights = anchors.new_zeros((num_anchors,))    # 回归权重 (107136,)
if self.norm_by_num_examples:   # False
    num_examples = (labels >= 0).sum()
    num_examples = num_examples if num_examples > 1.0 else 1.0
    reg_weights[labels > 0] = 1.0 / num_examples
else:
    reg_weights[labels > 0] = 1.0    # 将前景anchor的权重赋1

最后，将这构建的3个Tensor进行字典保存，返回到assign_targets函数中。

ret_dict = {
    
    
    'box_cls_labels': labels,           # 背景anchor的label是0，前景的anchor的label是当前处理的类别1  (107136,)
    'box_reg_targets': bbox_targets,    # 编码后待模型预测拟合的结果，背景anchor的编码信息也是0 (107136,7)
    'reg_weights': reg_weights,         # 背景anchor权重为0，前景anchor权重为1  (107136,)
}
return ret_dict

总结：本质上这个函数就是根据iou选择出前景anchor，然后对其进行类别赋值以及与其匹配的gt编码信息赋值，还设置了其回归的权重为1.

---