3.0 SA-SSD辅助网络详述

写在前面：本文主要结合源码，理解SA-SSD辅助网络的设计思想和实现原理。
涉及代码：single_stage.py中forward_train函数前四行代码。该四行代码为SA-SSD的核心代码，完成了论文中Backbone network和Auxiliary network
其他博客：
1.0 SA-SSD 环境配置
2.0 SA-SSD KITTI 3D数据可视化

3.0 SA-SSD辅助网络详述

    def forward_train(self, img, img_meta, **kwargs):
        # img [1, 3, 384, 1248]
        batch_size = len(img_meta)
        # step1: 处理多batch情况
        ret = self.merge_second_batch(kwargs)
        # step2: 
        vx = self.backbone(ret['voxels'], ret['num_points'])
        # step3:
        x, conv6, point_misc = self.neck(vx, ret['coordinates'], batch_size, is_test=False)
        ...

（一）输入说明

以单个batch size举例：

/	img	img_meta	kwargs
内容	[1, 3, 384, 1248]	数组：记录图像信息	字典：记录变量信息
说明	第一个维度表示batch size	数组长度是batch size。每一个数组元素是字典，包括标定参数、图像ID等	包括预选框、体素中点、真值框、真值标签等

举例说明img_meta

单个batch size时，img_meta 仅有一个元素：img_meta[0]

img_meta[0]	内容	说明
“calib”	包含P2、P3等多个参数信息	TODO 具体含义参考KITTI数据集
“img_shape”	(375, 1242, 3)	TODO 为什么与img的尺寸不一致？
“sample_idx”	3968	样本ID

举例说明kwargs

kwargs主要在kitti.py的prepare_train_img函数生成。

设置类别仅有一个Car

kwargs	内容	说明
“anchors”	数组 [70400, 7]	70400个Car预选框；7个参数分别是: x, y, z, w, h, d, cor
“anchor_mask”	数组 [70400]	TODO 值为True和False，不确定具体含义
“gt_bboxes”	数组 [13, 7]	13个真实框；7个参数分别是: x, y, z, w, h, d, cor
“gt_labels”	数组 [13]	13个真实框的标签ID，在这里只有一个类别0
“gt_types”	数组 [13]	13个真实框的标签名字，在这里只有一个类别Car
“coordinates”	数组 [17727, 3]	17727个体素的中点坐标
“num_points”	数组 [17727]	17727个体素的点云数量
“voxels”	数组 [17727, 5, 4]	每个体素至多有5个点云，每个点云有4个初始特征（x, y, z, 反射率）

（二）步骤一：

ret = self.merge_second_batch(kwargs)

• 目的：处理多batch的情况
• 返回值：ret
  一个字典，主要结构与kwargs类似。略有改变，改变均与 batch size 相关。
  例如，"coordinates"保存内容由 [17727, 3] 改为 [17727, 4] ，新增的第一个维度表示 batch size。

（三）步骤二：

vx = self.backbone(ret['voxels'], ret['num_points'])

• 输入说明
  ret['voxels'] [17727, 5, 4] 体素内所有点云的四维特征（不超过5个点云）
  ret['num_points'] [17727] 体素内的点云个数
• 目的
  基于每个体素，求其内部所有点云4维特征分别的平均值。即平均的x, y, z 坐标与平均反射率。
• 返回值说明
  vx [17727, 4]
  17727个体素，4维特征分别是平均的x, y, z 坐标与平均反射率。

（四）步骤三：【关键】

x, conv6, point_misc = self.neck(vx, ret['coordinates'], batch_size, is_test=False)

• 输入说明
  vx [17727, 4] 体素的平均4维特征（4维分别是x, y, z 坐标与反射率）
  ret['coordinates'] [17727, 4] 体素中心点的坐标 （4维分别是 batch size, x, y, z 坐标）
  batch_size
is_test 训练过程设置为False，开启辅助网络
• 目的
  实现论文中的Backbone network和Auxiliary network部分。注意，Backbone network也在步骤三中实现，而非步骤二。
• 返回值
  x
conv6
point_misc
  详见 （六）步骤三中Neck的具体实现

步骤三neck的具体实现并非十分简单，因此，将在（六）中详细说明。
在解释步骤三的具体实现前，先预备了解稀疏卷积的具体实现。

（五）预备知识：稀疏卷积

1.使用稀疏卷积的原因

（1）在正常卷积时，随着输入维度的升高，卷积计算量呈指数上升。

理解：

• 输入：M通道的2维图像
  使用N个2D卷积，计算量为$3^2\ast M\ast N$
• 输入：M通道的3维点云
  使用N个3D卷积，计算量为$3^3\ast M\ast N$

（2）如果特征在$d$维空间是稀疏的，没有必要在$3^d$空间遍历所有点。

2.稀疏卷积结果的具体表示

为了快速理解SA-SSA的辅助网络代码，这里只介绍稀疏卷积输入输出的表示方法，不涉及稀疏卷积的原理。
主要参考博客，如果想要了解稀疏卷积的原理同样可以参考上述博客

使用三个变量表示：

• $T$：稀疏特征图 $[W,H,D,m]$

理解：共有$W\ast H\ast D$个元素，每个元素有$m$维特征。其中，绝大部分元素的特征是空值。

• $M$：特征矩阵 $[a,m]$

理解：上述的$W\ast H\ast D$个元素中，共有$a$个非空元素。$a$个非空元素与其对应的$m$维特征，组成特征矩阵$M$。

• $H$：哈希表。Key是$M$中的行数，Value是$T$中的索引

理解：Key的取值范围是$[0,a-1]$，Vaule是一个三维向量$[w,h,d]$。利用该哈希表，可以知道特征矩阵中的每一个元素，在三维坐标中的实际位置。

注意：在已知特征矩阵$M$，哈希表$H$ 和 稀疏特征图$T$的尺寸后，可以得到稀疏特征图$T$

（六）步骤三中Neck的具体实现

self.neck()的核心代码见cmn.py中spMiddleFHD类的forward函数，主要分为两大部分：

1. 前一半，实现Backboe network
2. 后一半，实现Auxiliary network 【仅在训练阶段开启】

（1）Backboe network 实现

    def forward(self, voxel_features, coors, batch_size, is_test=False):
        points_mean = torch.zeros_like(voxel_features)      # points_mean 记录了batch id + 3维平均坐标
        points_mean[:, 0] = coors[:, 0]     # 获得batch id
        points_mean[:, 1:] = voxel_features[:, :3]      # 获得每个体素的前三个特征：3个点云的平均坐标

        coors = coors.int()
        x = spconv.SparseConvTensor(voxel_features, coors, self.sparse_shape, batch_size)       # 初始化SparseConvTensor
        x, middle = self.backbone(x)    # x: [5, 200, 176]  middle: [20, 800, 704] [10, 400, 352] [5, 200, 176]

        x = x.dense()   # [1, 64, 5, 200, 176] 第一个维度是batch size
        N, C, D, H, W = x.shape
        x = x.view(N, C * D, H, W)  # [1, 320, 200, 176]

        x, conv6 = self.fcn(x)  # x: [1, 256, 200, 176], conv6: [1, 256, 200, 176]

1. 输入说明

voxel_features [17727, 4] 体素的平均4维特征（4维分别是x, y, z 坐标与反射率）
coors [17727, 4] 体素中心点的坐标 （4维分别是 batch size, x, y, z 坐标）
batch_size
is_test 训练过程设置为False，开启辅助网络

2. 中间变量points_mean说明

points_mean [17727, 4] 4维分别是：batch size + 体素的平均3维特征（x, y, z 坐标）

实质上 points_mean 仅在 Auxiliary network 中使用，这里只是提前计算。

3. 稀疏卷积操作说明

步骤1）初始化类SparseConvTensor

相关代码

x = spconv.SparseConvTensor(voxel_features, coors, self.sparse_shape, batch_size)

目的

初始化x为SparseConvTensor类，该类具有以下属性：

• features: 值为voxel_features
• indices: 值为coors
• sparse_shape：值为sparse_shape
• batch_size：值为batch_size

原理

• voxel_features 相当于（五）中的$M$，即非空体素的特征。
• coors 相当于（五）中的$H$，即记录了非空体素的实际位置。
• sparse_shape 值为$[40,1600,1408]$，在配置文件car_cfg.py中设置。相当于（五）中$T$的尺寸，即$D,H,W$。

由此，根据以上参数，可以获得完整的稀疏特征图$T$信息。

步骤2）稀疏卷积操作

相关代码

x, middle = self.backbone(x)

目的

进行稀疏3D卷积，提取稀疏3D点云特征。完成Backboe network。

输出说明

• x 是一个SparseConvTensor类，是Backbone network的最终输出，其属性如下：

feature	indices	spatial_shape	batch_size
[12083, 64]	[12083, 4]	[5, 200, 176]	1

理解：最终提取了64维特征

• middle 是长度为3的数组，每个元素都是SparseConvTensor类，是Backbone network的中间输出。

4. 后续操作

相关代码

		x = x.dense()   # [1, 64, 5, 200, 176] 第一个维度是batch size
        N, C, D, H, W = x.shape
        x = x.view(N, C * D, H, W)  # [1, 320, 200, 176]

        x, conv6 = self.fcn(x)  # x: [1, 256, 200, 176], conv6: [1, 256, 200, 176]

目的

完成Detection network中的Reshape操作等，与辅助网络无关，不做说明。

（2）Auxiliary network 实现

        if is_test:
            return x, conv6
        else:
            # auxiliary network
            vx_feat, vx_nxyz = tensor2points(middle[0], (0, -40., -3.), voxel_size=(.1, .1, .2)) # vx_feat [34606, 32] vx_nxyz [34606, 4]
            p0 = nearest_neighbor_interpolate(points_mean, vx_nxyz, vx_feat)    # p0 [17727, 32]

            vx_feat, vx_nxyz = tensor2points(middle[1], (0, -40., -3.), voxel_size=(.2, .2, .4))
            p1 = nearest_neighbor_interpolate(points_mean, vx_nxyz, vx_feat)    # p1 [17727, 64]

            vx_feat, vx_nxyz = tensor2points(middle[2], (0, -40., -3.), voxel_size=(.4, .4, .8))
            p2 = nearest_neighbor_interpolate(points_mean, vx_nxyz, vx_feat)    # p2 [17727, 64]
            # 均为全连接层
            pointwise = self.point_fc(torch.cat([p0, p1, p2], dim=-1))      # pointwise [17727, 64]
            point_cls = self.point_cls(pointwise)                           # point_cls [17727, 1]
            point_reg = self.point_reg(pointwise)                           # point_cls [17727, 3]

            return x, conv6, (points_mean, point_cls, point_reg)

1. 输入说明

仅与Backbone network 的中间输出 middle 有关。
middle 是长度为3的数组，每个元素都是SparseConvTensor类。这里详细以middle[0]举例说明，其属性如下：

feature	indices	spatial_shape	batch_size
[34606, 32]	[34606, 4]	[20, 800, 704]	1

2. tensor2points 操作

相关代码

vx_feat, vx_nxyz = tensor2points(middle[0], (0, -40., -3.), voxel_size=(.1, .1, .2)) # vx_feat [34606, 32] vx_nxyz [34606, 4]

# transforms.py
def tensor2points(tensor, offset=(0., -40., -3.), voxel_size=(.05, .05, .1)):
    indices = tensor.indices.float()
    offset = torch.Tensor(offset).to(indices.device)
    voxel_size = torch.Tensor(voxel_size).to(indices.device)
    indices[:, 1:] = indices[:, [3, 2, 1]] * voxel_size + offset + .5 * voxel_size      # 得到实际 point_cloud_range [0, -40., -3., 70.4, 40., 1.] 中的位置
    return tensor.features, indices

参数说明

• tensor: middle[0] 其相关属性见上面的表格
• offset: $(0,-40,-3)$

1. KITTI数据集中，初始点云的point_range左下角点坐标即为$(0,-40,-3)$。
2. 该元素的目的：
   将映射后得到的点云坐标系$Q$进行平移，使其与初始点云坐标系$P$重合
   【$Q$中点云左下角点坐标为$(0,0,0)$，$P$中点云左下角点坐标为$(0,-40,-3)$】

• voxel_size: $(0.1,0.1,0.2)$

1. 初始点云的体素尺寸：$(0.05,0.05,0.1)$
2. 这里体素尺寸翻倍的原因：
   初始点云的体素结构是$[40,1600,1408]$，middle[0]中的spatial_shape 是$[20,800,704]$。
   即middle[0]中的一个体素，对应$2\ast 2\ast 2$个初始点云的体素。

主要目的

• 原先tensor.indices中后三维记录的是tensor.spatial_shape中的位置。
  【即体素的位置】
• 现在需要获得在原始点云 point_cloud_range [0, -40., -3., 70.4, 40., 1.] 中的位置。
  【即点云的位置】

映射操作

 indices[:, 1:] = indices[:, [3, 2, 1]] * voxel_size + offset + .5 * voxel_size  

• indices[:, [3, 2, 1]] * voxel_size：修改坐标系的尺度，获得所有体素左下角的坐标
• + offset：平移坐标系，使其与初始点云坐标系$P$完全重合
• + .5 * voxel_size：获得所有体素中点在初始点云坐标系$P$中的实际坐标

输出说明

• vx_feat [34606, 32] 即middle[0].feature
• vx_nxyz [34606, 4]
  第1维是batch_size
  第2-4维是上述34606个体素中点，在实际 point_cloud_range [0, -40., -3., 70.4, 40., 1.] 中的位置

3. nearest_neighbor_interpolate 操作

1）相关代码

p0 = nearest_neighbor_interpolate(points_mean, vx_nxyz, vx_feat)    # p0 [17727, 32]

# cmn.py
def nearest_neighbor_interpolate(unknown, known, known_feats):
    """
    :param pts: (n, 4) tensor of the bxyz positions of the unknown features         [17727, 4]
    :param ctr: (m, 4) tensor of the bxyz positions of the known features           [34606, 4]
    :param ctr_feats: (m, C) tensor of features to be propigated                    [34606, 32]
    :return:
        new_features: (n, C) tensor of the features of the unknown features
    """
    dist, idx = pointnet2_utils.three_nn(unknown, known)    # dist: [17727, 3], idx: [17727, 3]
    dist_recip = 1.0 / (dist + 1e-8)                        # [17727, 3]
    norm = torch.sum(dist_recip, dim=1, keepdim=True)       # [17727, 1]
    weight = dist_recip / norm                              # 权重值 [17727, 3]
    interpolated_feats = pointnet2_utils.three_interpolate(known_feats, idx, weight)        # [17727, 3]

    return interpolated_feats

2）参数说明

• unknown: points_mean $[17727,4]$
  在 Backboe network 中介绍过，17727表示初始非空体素的个数，4维分别是：batch size + 体素的平均3维特征（x, y, z 坐标）
• known: vx_nxyz $[34606,4]$

1. tensor2points的输出，记录已知特征的点云的实际位置。
2. 34606表示映射后拥有特征的体素个数，4维分别是：第1维是batch_size；第2-4维是上述34606个体素中点，在实际 point_cloud_range [0, -40., -3., 70.4, 40., 1.] 中的位置。

• known_feats: known_feats $[34606,32]$

1. tensor2points的输出，记录点云的特征。
2. 34606表示映射后拥有特征的体素个数，32维是特征。

3）主要目的

根据已知的坐标点及其特征，通过特征插值，获得目标坐标点的特征。

4）插值操作

4)-1. 论文中的插值操作：

• 参数说明
  $p_i$表示目标待求点，$p_j$表示$p_i$附近的已知点，$M$为已知点的数量。（此处，“附近”是一个球，不同尺度的输出有不同的半径）

• 主要思想
  针对每一个待求点，寻找其周围已知点，并分别计算距离权重。根据已知点的特征及其距离权重，计算得到待求点的特征。

4)-2. 代码实现中的插值操作：

主要思想：
针对每一个待求点，寻找最近的三个已知点。根据这三个点的特征值和距离权重，计算待求点的特征。

1.计算最近的三个已知点：

    dist, idx = pointnet2_utils.three_nn(unknown, known)    # dist: [17727, 3], idx: [17727, 3]

• dist: [17727, 3] 共有17727个未知点，3维表示3个最近已知点至该未知点的距离。
• idx: [17727, 3] 这里的3维记录了3个最近已知点的ID。
• 其他说明
  three_nn函数主要由interpolate_gpu.cu中的three_nn_kernel_fast函数实现。由于该函数涉及大量遍历计算，因此用C++编写，在环境配置阶段已经对该C++模块进行编译。

2.距离权重计算

    dist_recip = 1.0 / (dist + 1e-8)                        # [17727, 3]
    norm = torch.sum(dist_recip, dim=1, keepdim=True)       # [17727, 1]
    weight = dist_recip / norm                              # 权重值 [17727, 3]

针对每一个未知点（共有17727个），计算其对应三个已知点的权重，并进行归一化。

3.特征计算

    interpolated_feats = pointnet2_utils.three_interpolate(known_feats, idx, weight)        # [17727, 3]

• 参数说明
  1.known_feats [34606, 32] 共有34606个已知点，每个点有32维特征。
  2.idx [17727, 3] 共有17727个未知点，3维记录了3个最近已知点的ID，$ID\in [0,34605]$。结合known_feats，可以得到3个最近已知点的32维特征。
  3.weight [17727, 3] 记录了3个最近已知点的权重
• 函数说明
  three_interpolate函数主要由interpolate_gpu.cu中的three_interpolate_kernel_fast函数实现。主要计算过程就是针对每个未知点，分别计算特征的加权和。
• 输出说明
  interpolated_feats [17727, 3]

4. 特征拼接、预测、回归

相关代码

 # 均为全连接层
            pointwise = self.point_fc(torch.cat([p0, p1, p2], dim=-1))      # pointwise [17727, 64]
            point_cls = self.point_cls(pointwise)                           # point_cls [17727, 1]
            point_reg = self.point_reg(pointwise)                           # point_reg [17727, 3]

均为全连接操作，不做详细说明。

（3）Neck的输出

• x [1, 256, 200, 176]
  属于 Backboe network 的输出
• conv6 [1, 256, 200, 176]
  属于 Backboe network 的输出
• point_misc
  属于 Auxiliary network 的输出。包含：
  points_mean [17727, 4]
  point_cls [17727, 1]
  point_reg [17727, 3]

至此，Neck部分的实现已经阐述完毕。
single_stage.py中forward_train函数前四行代码的讲解也随之阐述完毕。
SA-SSD辅助网络的搭建均在Neck部分实现，至于训练时的损失值计算将在后续进行说明。

（七）对辅助网络的理解

体素法的缺点：

对体素内部点云结构的感知能力较弱。理由：体素的特征由多个内部点云的平均特征表示，较为粗粒度。【即平均x, y, z, 反射率】

辅助网络为什么能提高结构感知能力？

当体素内部结构变化时，一方面体素特征会改变，即最终提取的特征会改变【即 tensor 会改变】；另一方面unknown的点坐标会改变【即 point 坐标会改变】；则 tensor2points 操作后，unknown的点特征会改变。
应用这种辅助网络，额外添加监督任务，就是对这种改变进行额外的监督。因此能更好的提高结构感知能力。

辅助网络为什么提高的是特征提取网络的结构感知能力？

由于特征插值和拼接操作都是非学习的，特征插值的目标点能体现点云内部结构，也就是说得到的映射后的点云特征是由特征提取网络学习得到的，同时又进一步体现结构特征。通过辅助网络检测头，进行前景分割，进行附加的有监督学习，从而监督特征提取网络感知体素内部的点云特征。