OpenPCDet系列 | 5.4 PointPillars算法——AnchorHeadSingle模型预测头模块

OpenPCDet的整个结构图：

PointPillars算法属于OneStage算法，其中onestage可供选择的dense_head在__init__函数中选择，这里PointPillars算法使用的是AnchorHeadSingle，其又是继承自AnchorHeadTemplate。

# 根据MODEL中的DENSE_HEAD确定选择的模块
__all__ = {
    
    
    'AnchorHeadTemplate': AnchorHeadTemplate,   # 基类
    'AnchorHeadSingle': AnchorHeadSingle,
    'PointIntraPartOffsetHead': PointIntraPartOffsetHead,
    'PointHeadSimple': PointHeadSimple,
    'PointHeadBox': PointHeadBox,
    'AnchorHeadMulti': AnchorHeadMulti,
    'CenterHead': CenterHead
}

所以，下面先记录一下AnchorHeadTemplate的作用，再记录AnchorHeadSingle模块

---

AnchorHeadTemplate模块

由于PointPillars是属于OneStage的anchor-based算法，所以首先就涉及到了anchor的生成，这部分是通过AnchorGenerator.generate_anchors类函数来实现。此外，在AnchorHead中还需要定义回归的编码方式以及目标对齐方式，最后再构建分类头、回归头、方向预测头的损失函数。

# 功能：dense head模块的基类
class AnchorHeadTemplate(nn.Module):
    def __init__(self, model_cfg, num_class, class_names, grid_size, point_cloud_range, predict_boxes_when_training):
        """
        Args:
            model_cfg:      DENSE_HEAD的配置文件
            num_class:      类别数目(3类)
            class_names:    类别名称: ['Car', 'Pedestrian', 'Cyclist']
            grid_size:      网格大小
            point_cloud_range:  点云范围：[-x, -y, -z, x, y, z]
            predict_boxes_when_training:    布尔变量：False (twoStage模型才会设置为True)
        """
        super().__init__()    # 初始化nn.Module
        self.model_cfg = model_cfg
        self.num_class = num_class
        self.class_names = class_names
        self.predict_boxes_when_training = predict_boxes_when_training  # False (twoStage模型才会设置为True)
        self.use_multihead = self.model_cfg.get('USE_MULTIHEAD', False) # False (多尺度head的设置)

        # Dense Head模块包含三大子部分：
        # 1)对生成的anchor和gt进行编码和解码
        anchor_target_cfg = self.model_cfg.TARGET_ASSIGNER_CONFIG   # anchor分配文件
        self.box_coder = getattr(box_coder_utils, anchor_target_cfg.BOX_CODER)( # 在box_coder_utils文件中调用ResidualCoder类
            num_dir_bins=anchor_target_cfg.get('NUM_DIR_BINS', 6),  # 如果没有设置，默认为6
            **anchor_target_cfg.get('BOX_CODER_CONFIG', {
    
    })
        )

        # 2)anchor生成配置
        anchor_generator_cfg = self.model_cfg.ANCHOR_GENERATOR_CONFIG   # list:存储每个类别的anchor生成设置
        anchors, self.num_anchors_per_location = self.generate_anchors(
            anchor_generator_cfg, grid_size=grid_size, point_cloud_range=point_cloud_range,
            anchor_ndim=self.box_coder.code_size
        )

        # 3)gt匹配
        self.anchors = [x.cuda() for x in anchors]  # 放在GPU上
        self.target_assigner = self.get_target_assigner(anchor_target_cfg)

        # 4)保存前向传播结果并构建损失函数
        self.forward_ret_dict = {
    
    }    # 根据forward_ret_dict内容来计算loss
        self.build_losses(self.model_cfg.LOSS_CONFIG)       # 分类损失、回归损失、方向损失的构建
        
        ......

1. AnchorGenerator

初始化参数变量如下所示：

函数的最后返回的是anchor_list列表以及每个位置每个类别有多少种anchor的列表，如下所示：

2. ResidualCoder

初始化参数变量如下：

3. AxisAlignedTargetAssigner

初始化参数变量如下：

---

AnchorHeadSingle模块

在PointPillars算法中，dense_head模块是最难的部分，其初始化和前向传播都涉及多个部分，下面分别进行介绍。

1. AnchorHeadSingle初始化

AnchorHeadSingle模块初始化的结构图：

在配置文件中这里选择的预测3个类别：[‘Car’, ‘Pedestrian’, ‘Cyclist’]，每个类别的anchor存在两个方向，也就是一共会生成6种anchor。对于每个anchor需要7个维度的信息表示：[x, y, z, dx, dy, dz, heading]。对应对于分类head的输出维度是6 * 3类别=18；对于回归reg head的输出维度是6 * 7个信息表示=42；最后还有分类6 * 2个方向=12。同时，不同的head设置了不同的损失函数。

模型的具体构建结果如下所示，每个head只有一层的linear来进行最后的预测。到了AnchorHeadSingle模块中，具体的定义其实只涉及到了模型的定义。（其中损失的定义还是在基类中进行构建的）

AnchorHeadSingle(
  (cls_loss_func): SigmoidFocalClassificationLoss()
  (reg_loss_func): WeightedSmoothL1Loss()
  (dir_loss_func): WeightedCrossEntropyLoss()
  (conv_cls): Conv2d(384, 18, kernel_size=(1, 1), stride=(1, 1))
  (conv_box): Conv2d(384, 42, kernel_size=(1, 1), stride=(1, 1))
  (conv_dir_cls): Conv2d(384, 12, kernel_size=(1, 1), stride=(1, 1))
)

2. AnchorHeadSingle训练前向传播

AnchorHeadSingle模块前向传播的结构图如下所示：

传入dense_head时的数据字典结构如下所示，这里主要需要对spatial_feature_2d的特征进行进一步的处理：

在AnchorHeadSingle模块时已经分别构建的各分类、回归的head，那么这里首先就是让spatial_feature_2d特征分别通过这些head来进行对应object的预测，然后见所有的预测结构存储在self.forward_ret_dict字典中。

随后根据相关的gt boxes信息构建gt信息，调用基类的self.target_assigner.assign_targets函数，同样的将结果存储在self.forward_ret_dict字典中

后续的过程就是机遇这个forward_ret_dict字典，分别调用self.get_cls_layer_loss()、self.get_box_reg_layer_loss()来进行具体的损失计算。

3. AnchorHeadSingle测试前向传播

在还没进行是训练还是测试过程时候，此时的data_dict还是一致的。

但是，在测试过程中，这里根据的各预测的特征矩阵来生成box，代码如下，更新了是两个预测信息：batch_cls_preds和batch_box_preds。

# 测试过程
if not self.training or self.predict_boxes_when_training:
    batch_cls_preds, batch_box_preds = self.generate_predicted_boxes(   # 根据各类预测矩阵生成预测box
        batch_size=data_dict['batch_size'],     # 设置的batch_size为16
        cls_preds=cls_preds, box_preds=box_preds, dir_cls_preds=dir_cls_preds   # 各预测特征map
    )
    data_dict['batch_cls_preds'] = batch_cls_preds  # (16, 321408, 7)
    data_dict['batch_box_preds'] = batch_box_preds  # (16, 321408, 3)
    data_dict['cls_preds_normalized'] = False

将预测的特征矩阵更新在data_dict中，返回的data_dict结果如下所示：

随后，返回结束了dense_head模块的处理返回的pointpillars算法中，在后续的处理是将更新后的data_dict传入到检测模型的基类Detector3DTemplate.post_processing函数中进行后处理操作。

---