OpenPCDet series | 5.4 PointPillars algorithm - AnchorHeadSingle model prediction head module

The entire structure diagram of OpenPCDet:

The PointPillars algorithm belongs to the OneStage algorithm, where the dense_head that onestage can choose is selected in the __init__ function. Here, the PointPillars algorithm uses AnchorHeadSingle, which is inherited from AnchorHeadTemplate.

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

So, let's first record the function of AnchorHeadTemplate, and then record the AnchorHeadSingle module

---

AnchorHeadTemplate module

Since PointPillars is an anchor-based algorithm belonging to OneStage, it first involves the generation of anchors, which is AnchorGenerator.generate_anchorsrealized through class functions. In addition, in AnchorHead, it is also necessary to define the regression encoding method and target alignment method, and finally construct the loss function of the classification head, regression head, and direction prediction head.

# 功能：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

Initialize parameter variables as follows:

At the end of the function, the anchor_list list and the list of how many kinds of anchors there are in each category of each position are returned, as follows:

2. ResidualCoder

Initialize parameter variables as follows:

3. AxisAlignedTargetAssigner

Initialize parameter variables as follows:

---

AnchorHeadSingle module

In the PointPillars algorithm, the dense_head module is the most difficult part. Its initialization and forward propagation involve multiple parts, which will be introduced separately below.

1. AnchorHeadSingle initialization

The structure diagram of AnchorHeadSingle module initialization:

In the configuration file, the predictions selected here are 3 categories: ['Car', 'Pedestrian', 'Cyclist']. There are two directions for the anchor of each category, that is, a total of 6 anchors will be generated . For each anchor, information representation of 7 dimensions is required: [x, y, z, dx, dy, dz, heading]. The output dimension corresponding to the classification head is 6 * 3 categories = 18; the output dimension for the regression reg head is 6 * 7 information representation = 42; finally there are classification 6 * 2 directions = 12 . At the same time, different heads set different loss functions.

The specific construction results of the model are shown below. Each head has only one layer of linear for final prediction. In the AnchorHeadSingle module, the specific definition actually only involves the definition of the model. (The definition of loss is still constructed in the base class)

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 training forward propagation

The structure diagram of the forward propagation of the AnchorHeadSingle module is as follows:

The data dictionary structure when passing in dense_head is as follows, here it is mainly necessary to further process the features of spatial_feature_2d:

In the AnchorHeadSingle module, the head of each classification and regression has been constructed separately, so the first thing here is to let the spatial_feature_2d feature predict the corresponding object through these heads, and then see that all the prediction structures are stored in the self.forward_ret_dict dictionary.

Then build gt information based on the relevant gt boxes information, call the self.target_assigner.assign_targets function of the base class, and store the result in the self.forward_ret_dict dictionary in the same way

The subsequent process is to use the forward_ret_dict dictionary, and call self.get_cls_layer_loss() and self.get_box_reg_layer_loss() to perform specific loss calculations.

3. AnchorHeadSingle test forward propagation

When the training or testing process has not yet been performed, the data_dict at this time is still consistent.

However, during the test, the box is generated based on the feature matrix of each prediction. The code is as follows, and two prediction information are updated: batch_cls_preds and 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

Update the predicted feature matrix in data_dict, and the returned data_dict result is as follows:

Subsequently, in the pointpillars algorithm returned after the processing of the dense_head module is returned, the subsequent processing is to pass the updated data_dict into the base class function of the detection model for post-processing Detector3DTemplate.post_processingoperations.

---