【代码阅读】Part-A^2 Net 代码

前一段时间Part-A2的代码公布了，这里就记录一下自己的学习过程。
github代码仓库：传送门
对文章的解读：传送门

安装

其中，spare-conv安装起来有点难度。具体可以参考一篇SECOND的安装的博客，相比于在spare-conv的github中的安装介绍，我采坑的一点是要安装：
sudo apt-get install libboost-all-dev
否则会编译失败。还有就是scikit-image安装时出现错误，更新pip也许可以解决，参考另外一篇博客。

数据预处理

要求运行一下：

python kitti_dataset.py create_kitti_infos

这就看一下这个干了些什么：

# pcdet\datasets\kitti\kitti_dataset.py
def create_kitti_infos(data_path, save_path, workers=4):
    ...
    print('---------------Start to generate data infos---------------')

    dataset.set_split(train_split)
    kitti_infos_train = dataset.get_infos(num_workers=workers, has_label=True, count_inside_pts=True)
    with open(train_filename, 'wb') as f:
        pickle.dump(kitti_infos_train, f)
    print('Kitti info train file is saved to %s' % train_filename)
    ...
    print('---------------Start create groundtruth database for data augmentation---------------')
    dataset.set_split(train_split)
    dataset.create_groundtruth_database(train_filename, split=train_split)

    print('---------------Data preparation Done---------------')

可以看到这个函数干了两个事情：
1、准备了train和val的infos，也就dataset.get_infos这个函数。具体的可以看一下这个函数内部，这里就简单说一下，这个函数内部把train和val对应的所有的label都拿出来，然后去掉不相关的（例如 DonotCare类别），最重要的就是把在相机坐标系下的3D bouding boxes的lable转到了LiDAR坐标系下。然而这个也是只转了中心点的坐标xyz，长宽高与坐标系无关，所以不用转，ry也没转，这个操作我不太理解，为什么不统一直接转过去好了。
2、准备gt_database，将train下的要预测的box都拿出来，并且把其中的点都抠出来，这一步在gt_aug中会用到，这个在PointRCNN中也出现过。

dataset

按照我的习惯，我会在看train之前看一下dataset，是怎么做get_item的，否则，train里面的变量都搞不清楚。

dataset在train中是直接调用build_dataloader的

# -----------------------create dataloader & network & optimizer---------------------------
    train_set, train_loader, train_sampler = build_dataloader(
        cfg.DATA_CONFIG.DATA_DIR, args.batch_size, dist_train, workers=args.workers, logger=logger, training=True
    )

而这个build_dataloader是在pcdet/dataset/__init__.py中定义的。

代码中一共有3个跟dataset有关的类，继承关系如下：
DatasetTemplate$\to$BaseKittiDataset$\to$KittiDataset
三个类的作用可以理解为如下：

• DatasetTemplate：是做一些对所有数据集都要做的一些工作，例如数据增广等。PartA2中生成voxel的mask的这个target的准备工作也是在这个类中完成的
• BaseKittiDataset：做一些加载KITTI数据有关的类，包括加载point和label
• KittiDataset：是直接产生训练数据的，调用父类的一些函数

这里就跟着KittiDataset看，其他两个类中的一些辅助函数就忽略了。直接看KittiDataset的__getitem__函数，就是在加载point和label，label是实在预处理中处理过的。然后就调用了父类DatasetTemplate中的prepare_data函数，这个函数是准备data中最重要的，这里就不放代码了，说一下这个函数干了啥，主要是说training的准备过程：

1. 先使用road_plan和预处理中做好的gt_database，做gt_aug，就是把其他场景中的物体放到本场景中
2. 然后做per_object的旋转和加噪声点
3. 加入全局的旋转和尺度缩放
4. 使用voxel_generator将point变成voxel
5. 生成target，重要的是使用generate_voxel_part_targets函数，得到每个voxel的cls和reg的label

Training

看完了dataset，再回到model，看看model是怎么构建的。model的类也是有一个基类，其他类是其子类：
Detector3D$\to$PartA2Net
Detector3D$\to$PointPillar
Detector3D$\to$SECONDNet
这里只详细看PartA2Net

train.py中的构建，是通过调用build_network函数，这个定义在pcdet/dataset/__models__.py中。而这个函数的传入值定义在pcdet/tools/cfgs/PartA2.yaml中。

构造网络的过程，从初始化看起，那么第一步就是父类Detector3D.__init__()这个函数，然后就是PartA2Net.__init__()，可以看到，PartA2Net.__init__()调用了父类的build_networks()函数，这里就结合PartA2.yaml看一下model是如何构造的，主要分为如下几个模块：

• vfe：用的是MeanVoxelFeatureExtractor，在pcdet/models/vfe/vfe_utils.py中可以找到，可以看到，这个就是把每个voxel内部的多个点的feature平均取值为一个。
• rpn：用的是UNetV2，在pcdet/models/rpn/rpn_unet.py中，这个类的父类是UNetHead，也在同一个py中。用3D sparse conv定义了一个U-Net，包括U-Net的主体（self.convx,self.conv_up_x）、在金字塔尖的一个特征提取头(self.conv_out)、与U-Net输出链接的cls和reg头(self.seg_cls_layer,self.seg_reg_layer)。
• rpn_head：用的是UNetHead，在pcdet/models/bbox_head/rpn_head.py中可以找到，父类是AnchorHead。用3D sparse conv定义了多sacle特征提取模块（self.blocks，self.deblocks）和预测box相关的模块。
• rcnn：用的是SpConvRCNN，在pcdet/models/rcnn/partA2_rcnn_net.py中，继承父类RCNNHead。

Forward Inference

上面已经讲了model的构建过程，其实前向计算就是把上面的model各个子网络走一遍。train.py中调用trian_model函数（pcdet/tools/train_utils/train_utils.py中），然后在调用train_one_epoch，然后调用model_func来计算每个batch。从train.py中查看，可以找到，这个model_func是定义在pcdet/models/__init__.py中。可以看到，核心代码就是

ret_dict, tb_dict, disp_dict = model(input_dict)

这么，就可以定位到，PartA2这个网络的forward函数中了。可以看到，这个函数分别调用了forward_rpn，forward_rcnn，get_training_loss。接下来就一个一个看。接下来，各个部分的tensor均是batch=2。

forward_rpn

输入：
voxels：32000x5x4，其中32000是2个point cloud的voxel的总数，定义在PartA2.yaml中DATA_CONFIG:TRAIN:MAX_NUMBER_OF_VOXELS: 16000；5是代表一个voxel中最多包含5个点，定义在DATA_CONFIG:VOXEL_GENERATOR: MAX_POINTS_PER_VOXEL: 5；

预处理：

voxel_features = self.vfe(
    features=voxels,
    num_voxels=num_points,
    coords=coordinates
)
# voxel_feature: 32000x4

input_sp_tensor = spconv.SparseConvTensor(
    features=voxel_features,
    indices=coordinates,
    spatial_shape=self.sparse_shape,
    batch_size=batch_size
)
# input_sp_tensor：41x1600x1480，把voxel_feature变为sparse tensor

U-Net，这个在UNetV2.forward()中，这块的注释已经非常好了，这里就说一下其中计算过程：
1）通过encoder，将特征提取到金字塔顶，得到x_conv4，大小为[200, 176, 5]
2）对x_conv4进一步压缩特征，得到out，大小为[200, 176, 2]，然后合并代表高度的2，其中一个高度channel有128个通道，所以合并完高度的tensor为[2, 256, 200, 176]，第一个2为batch。这其实就是一个俯视图的feature map，与图像卷积出来的金字塔顶的map是一样的格式，都是batchchannelH*W。把这个feature叫做spatial_features
3）对x_conv4做decoder，得到与输入一样多数量的voxel
4）对每个voxel的特征计算一个cls和reg分量，分别是(32000x1, 32000x3)，对应于32000个voxel，这两个量分别被叫做u_seg_preds和u_reg_preds
5）加载target，这个dataset中已经讲过

rpn_head，这个在RPNV2.forward()中，输入是之前得到的spatial_features：
1）提取多尺度的特征，得到x_in，大小是[2, 512, 200, 176]
2）然后通过分类头和回归头得到cls_preds和box_preds，大小分别为[2, 18, 200, 176], [2, 42, 200, 176]，和dir_cls_preds，大小为[2, 200, 176, 12]。cls_preds中的18=63，6个anchor，3个class，目的是找出每个anchor生成的proposal对应的class，这个通过选择与box的semantic class有关。box_preds中的42=67，每个anchor生成proposal需要7个参数。dir_cls_preds中的12=62，2是代表这个anchor生成proposal时的正反方向。
3）通过assign_target得到target，就是把要预测的真值，与feautre map的grid中的anchor关联起来。这个与图像上的关联操作是一样的。每个grid一共有3类2方向的6个anchor，根据类别和IoU将真值与anchor关联。

forward_rcnn

有了rpn的输出，接着就进入forward_rcnn了。

先decode出proposal：
先通过box_preds，cls_preds，dir_cls_preds和anchor把rpn预测的proposal decode出来。anchor的大小是[2, 211200]，其中211200=1762006，6是指3类*2方向的anchor。这个过程是self.rpn_head.box_coder.decode_with_head_direction_torch，在pcdet/utils/box_coder_utils.py可以找到。具体就是通过box_preds和dir_cls_preds把anchor转为target。我认为dir_cls_preds的作用是把每个anchor二分类，分别指向正反两个方向。最后得到

proposal_layer
在pcdet/models/model_utils/proposal_layer.py中可以找到，具体操作就是把proposal按照score排序，取前9000个，NMS，然后取前512个proposal作为roi

SpConvRCNN.forward()：
1）找到roi对应的target，用proposal_target_layer选取64个前景roi和64个背景roi，计算rcnn的[x, y, z, w, h, l, ry]的target
2）roi pooling，得到roi-wise feature
3）卷积提取特征
4）使用分类头和回归头计算得到rcnn_cls和rcnn_reg，分别是[256, 1]和[256, 7]，256=2*128，batch是2，一帧里出128个box

get_training_loss

这个代码和文章中写的都比较清晰了，就不再赘述。

Testing

testing的前向计算过程与training一样，没有了loss的计算，多了由rpn和rcnn的输出得到预测结果的部分。

在testing中，proposal_layer出来的roi就是100个。从rcnn的结果中decode出来box，得到batch_cls_preds和batch_box_preds，大小分别为[2, 100]和[2, 100, 7]。

然后对每个batch中，根据batch_cls_preds排序，做NMS，得到最后的box。
box的semantic class由rpn的结果通过选择给出。