OpenPCDet初级教程【自定义模型、loss】

        最近在研究点云物体检测，基于OpenPCDet框架进行算法开发可以节约大量的重复性工作，专心集中在核心算法的设计上，大量节约时间。同时，因为框架由大公司专业团队进行维护，代码质量稳定。本文以小项目的形式记录如何采用OpenPCDet自定义模块、网络、loss实现训练。

目标任务：

        以OpenPCDet中的pointrcnn为基础，单独剥离其中的前景点分割网络，将POINT_HEAD模块替换为自定义的CLS_HEAD。采用新建detector、新建模块、新建网络的方式实现任务。目标网络简单：pointnet+【256 256】的mlp。

步骤一：网络构建

        在OpenPCDet中，有八个基本模块， 'vfe', 'backbone_3d', 'map_to_bev_module', 'pfe', 'backbone_2d', 'dense_head', 'point_head', 'roi_head'每个模块中都有若干网络可供选择。

①自定义detector：新建detector.py 若需要在8个基础模块外添加模块，把新建模块名加入module_topology簇

class PointRCNN_cls(Detector3DTemplate):
    def __init__(self, model_cfg, num_class, dataset):
        super().__init__(model_cfg=model_cfg, num_class=num_class, dataset=dataset)
        # 在框架中增加额外模块cls_head
        self.module_topology = [
            'vfe', 'backbone_3d', 'cls_head', 'map_to_bev_module', 'pfe',
            'backbone_2d', 'dense_head', 'point_head', 'roi_head'
        ]
        self.module_list = self.build_networks()

把新detector类加入簇，detector init

__all__ = {
    'Detector3DTemplate': Detector3DTemplate,
    'SECONDNet': SECONDNet,
    'PartA2Net': PartA2Net,
    'PVRCNN': PVRCNN,
    'PointPillar': PointPillar,
    'PointRCNN': PointRCNN,
    'SECONDNetIoU': SECONDNetIoU,
    'CaDDN': CaDDN,
    'VoxelRCNN': VoxelRCNN,
    'CenterPoint': CenterPoint,
    'PVRCNNPlusPlus': PVRCNNPlusPlus,
    'PointRCNN_cls': PointRCNN_cls,  # 自定义detector
}

②自定义模块：在detector3d_template定义新模块

def build_cls_head(self, model_info_dict):
    if self.model_cfg.get('CLS_HEAD', None) is None:
        return None, model_info_dict
    num_point_features = model_info_dict['num_point_features']
    # 从yaml读取网络相关配置
    cls_head_module = dense_heads.__all__[self.model_cfg.CLS_HEAD.NAME](
        model_cfg=self.model_cfg.CLS_HEAD,  # 根据NAME读取配置
        input_channels=num_point_features,
        num_class=self.num_class if not self.model_cfg.CLS_HEAD.CLASS_AGNOSTIC else 1,  # class-agnostic方式只回归2类bounding box，即前景和背景
    )
    model_info_dict['module_list'].append(cls_head_module)
    return cls_head_module, model_info_dict

③自定义网络：新建cls_head.py定义模块内的具体网络，这里采用两层mlp，网络的预测结果要存在batch_dict，方便模块化。

class ClsHead(CLS2_HeadTemplate):  # 从PointHeadTemplate调用make_fc_layers
    def __init__(self, num_class, input_channels, model_cfg, **kwargs):
    super().__init__(model_cfg=model_cfg, num_class=num_class)
    # 根据yaml构建网络
    self.cls_layers = self.make_fc_layers(
        fc_cfg=self.model_cfg.CLS_FC,
        input_channels=input_channels,
        output_channels=num_class
    )
    def forward(self, batch_dict):
    point_features = batch_dict['point_features']  # 从字典中获取每个点的特征 shape （batch * 16384, 128）
    point_cls_preds = self.cls_layers(point_features)  # (total_points, num_class)
    # 从每个点的分类预测结果中，取出类别预测概率最大的结果  (batch * 16384, num_class) --> (batch * 16384, )
    point_cls_preds_max, _ = point_cls_preds.max(dim=-1)
    # 将类别预测分数经过sigmoid激活后放入字典中
    batch_dict['point_cls_scores'] = torch.sigmoid(point_cls_preds_max)
    # 将点的类别预测结果和回归结果放入字典中
    ret_dict = {
        'point_cls_preds': point_cls_preds,
    }
    # 如果在训练模式下，需要根据GTBox来生成对应的前背景点，用于点云的前背景分割，给后面计算前背景分类loss
    if self.training:
        targets_dict = self.assign_targets(batch_dict)
        # 将一个batch中所有点的GT类别结果放入字典中 shape (batch * 16384)
        ret_dict['point_cls_labels'] = targets_dict['point_cls_labels']
    # 第一阶段生成的预测结果放入前向传播字典
    self.forward_ret_dict = ret_dict

    return batch_dict

把新网络加入模块的簇

__all__ = {
    'AnchorHeadTemplate': AnchorHeadTemplate,
    'AnchorHeadSingle': AnchorHeadSingle,
    'PointIntraPartOffsetHead': PointIntraPartOffsetHead,
    'PointHeadSimple': PointHeadSimple,
    'PointHeadBox': PointHeadBox,
    'AnchorHeadMulti': AnchorHeadMulti,
    'CenterHead': CenterHead,
    'ClsHead': ClsHead, # 自定义网络
}

④新建detector的yaml配置文件

注意YAML中的NAME和簇里的名字需要保持一致，也就是和类名保持一致

# 模型配置
MODEL:
    NAME: PointRCNN_cls  # 和detector定义中的名称保持一致：class PointRCNN_cls(Detector3DTemplate):

    BACKBONE_3D:
        NAME: PointNet2MSG
        SA_CONFIG:
            NPOINTS: [4096, 1024, 256, 64]
            RADIUS: [[0.1, 0.5], [0.5, 1.0], [1.0, 2.0], [2.0, 4.0]]
            NSAMPLE: [[16, 32], [16, 32], [16, 32], [16, 32]]
            MLPS: [[[16, 16, 32], [32, 32, 64]],
                   [[64, 64, 128], [64, 96, 128]],
                   [[128, 196, 256], [128, 196, 256]],
                   [[256, 256, 512], [256, 384, 512]]]
        FP_MLPS: [[128, 128], [256, 256], [512, 512], [512, 512]]

    CLS_HEAD:   # 模块名
        NAME: ClsHead  # 模型和定义的名字保持一致   class ClsHead(CLS2_HeadTemplate)
        CLS_FC: [ 256, 256 ]
        REG_FC: [ 256, 256 ]
        CLASS_AGNOSTIC: False
        USE_POINT_FEATURES_BEFORE_FUSION: False

        LOSS_CONFIG:

此时网络已构建完毕，可以通过detector里的self.module_list = self.build_networks()查看网络。

步骤二：loss构建

具体的loss如何计算在cls2_head_template中定义，get_cls_layer_loss作为CLS2_HeadTemplate的成员函数，通过self.cls_head.get_loss()调用，get_training_loss的作用是整理一下不同模块的loss，如pointrcnn里第一阶段loss和第二阶段loss。

①在detector中写forward，预测结果存在batch_dict里。

②在cls2 head template中修改assign_stack_targets，去除box估计的部分，只考虑分类。在detector中写get loss，整合一下不同网络的loss，如第一阶段里的cls和box的两部分loss。

def forward(self, batch_dict):
    for cur_module in self.module_list:  # 循环model list进行网络预测：pointnet-pointheadbox-pointrcnnhead
        batch_dict = cur_module(batch_dict)

    if self.training:
        loss, tb_dict, disp_dict = self.get_training_loss()

        ret_dict = {
            'loss': loss
        }
        return ret_dict, tb_dict, disp_dict
    else:
        pred_dicts, recall_dicts = self.post_processing(batch_dict)
        return pred_dicts, recall_dicts

def get_training_loss(self):
    disp_dict = {}
    loss_point, tb_dict = self.point_head.get_loss()

    loss = loss_point
    return loss, tb_dict, disp_dict

至此，自定义内部完毕，可以开始训练咯！！

python train.py + 自定义的yaml

参考：

OpenPCDet解读 - 知乎

OpenPCDet入门 - 知乎