八. 实战：CUDA-BEVFusion部署分析-spconv原理

前言

自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习下课程第八章——实战：CUDA-BEVFusion部署分析，一起来学习 spconv 原理

Note：之前在学习杜老师的课程中有简单记录过 Sparse Convolution 的一些基础知识，感兴趣的可以看下：复杂onnx解决方案（以sparseconv为例）

课程大纲可以看下面的思维导图

0. 简述

本小节的目标：理解 spconv 与普通的 conv 的区别，计算原理与流程，以及导出 onnx 需要考虑的事情

这节课给大家讲解第八章第 2 小节学习 spconv 的原理，通过这一小节我们去理解一下我们用 3D Sparse Convolution 对稀疏性的点云做卷积的时候我们应该怎么去做，spconv 的算法和流程是什么样的，它和普通卷积的区别有哪些，这是我们接下来需要去理解的。

最后我们作为第 3 小节的前述部分，会去给大家讲一下导出带有 Sparse Convolution 网络的 ONNX 的时候需要注意哪些事情

OK，我们下面正式开始

1. 举例分析spconv的计算流程

首先为了理解 spconv 的流程，我们把这个问题简单化，先看一个最简单的例子，我们假设现在有一个输入，大小是 1x5x5，同时经过一个 1x1x3x3 大小的卷积核得到 1x3x3 大小的输出，如下图所示：

卷积的整个过程 stride 是 1，padding 是 0，同时我们假设输入的 1x5x5 大小的数据中除了 in0 点以外其他的所有数据都为 0，也就是说只有 in0 点它是有值的，那么通过一个卷积之后我们知道输出的数据只有 out0 和 out1 有值，其他输出的数据全为 0

in0 跟 k1 计算得到 out0，in0 和 k0 计算得到 out1，卷积核每进行一次滑动计算，是需要做 9 次乘加法的，此外这个卷积核滑动完整个输入是需要 9 次，那也就意味着我们一共要做 9x9=81 次乘加法

那 81 次乘加法中其实只有 2 次乘加法是有必要的，那剩余的 79 次的计算其实都是在跟 0 做计算，都是没有必要的计算，我们其实可以 skip 掉

OK，那么我们要想办法把这两个计算给保存下来，我们先把这两个计算叫 atomic operation 也就是两个原子操作，那我们想要去保存这两个原子操作的话，我们需要怎么办呢

首先我们得需要知道 input 中 25 个数据哪个点它是有意义的，我们需要一个 table 来保存它，那么在 spconv 中我们会用 hash table 的方式去把 in0 这个点的坐标和它的一个编号给保存下来，也就是图中的 P_in

同理 output 也是一样的，它其实是只有两个数据是有意义的，那么把这两个数据也用 hash table 的方式给保存下来，保存它的一个坐标和编号，也就是图中的 P_out

OK，这两个 table 就完成了，那下一步要做什么呢，下一步就是说让两个 table 的各个元素给做出一个关联性，也就是说 input table 中的哪一个值和 output 中的哪个值是一一对应的，同时它这个对应的所需要的 kernel 的值是哪一个

那么这个对应关系我们可以用一个叫 Rulebook 的方式去保存，图中 Rulebook 有四列，分别代表的含义如下：

• (i,j)：代表的是 Kernel 的相对坐标，它相对坐标是以中间这个点为 (0,0) 原点的，来计算每一个点距离中心点的一个相对位置，如下图所示。比如 K0 这个点它距离中心点 (0,0) 的位置是 (-1,-1)，K1 这个点它距离中心点 (0,0) 的位置就是 (0,-1)
• count：代表的是编号，表示是这个点它要涉及的计算它的编号是第几个，那比如说 K0 它所涉及到的计算只有一个，那就是 count 等于 0，如果 K0 还涉及到另一个计算的话，那这个 count 就是 1
• v_in、v_out：代表着输入输出中非零点的对应的 index

对于 Rulebook 我们可以这么理解，我们先看第一行，v_in 中 index 等于 0 这个点也就是位置坐标是 (0,1) 的这个点（即 in0），它和 Kernel 中 (-1,-1) 这个点（即 K0）计算得到 v_out1（即 out1），同理第二行就是 v_in 也就是 in0 它跟需要跟 Kernel 的 (0,-1) 位置（即 K1）做计算 v_out0（即 out0），这样就可以把输入、输出和 Kernel 的计算关系给保存下来。

为了方便说明，图中我们也进行了举例说明，比如 in0 的值是 0.3，K0，K1 分别就是 0.4 和 0.9，最终得到的输出 out1 和 out0 分别是 0.27 和 0.12，那所以通过 rulebook 就可以把这个对应关系给绑定起来，这个是比较简单的例子

那下一步我们把这个问题再稍微复杂化一点，我们之前一直讨论的是输入 channel 等于 1，那么我们给它扩充一下把它变成三个维度，那么 input 就是 3x5x5，意味着 Kernel 它是 1x3x3x3，output 是没有发生变化的，因为 Kernel 数量没变，整个过程就如下图所示：

那么其实这个过程中虽然 input 和 kernel 的 channel 发生了变化，但是它坐标和坐标间的关系并没有改变，还是该哪个坐标跟哪个坐标乘，这个是不会发生变化的。那变化在哪呢，之前是一个元素和一个元素之间相乘得到一个元素，现在变成一个 vector 和一个 vector 相乘得到一个元素

那么理解这个之后我们再进一步扩展，下一步我们来扩展 kernel，那我们之前 Kernel 不是只有一个吗，现在我们扩充成两个，那么相应的 output 的 channel 也会变成 2，如下图所示：

那么在这种情况下它的坐标对应关系会发生改变吗，那也不会发生改变，哪个坐标跟哪个坐标相乘得到哪个坐标，这个关系是不会变的。那么变得是什么呢，变的是 kernel 它现在已经不再是一个 vector 而是一个 matrix 了，那么输出就是一个 vector 跟一个 Matrix 相乘

到目前为止我们一直聊的都是一个点，就是 input 中只有一个点做计算，那现在我们把问题再给它复杂化一点，input 中如果有两个点或者多个点时的情形，如下图所示：

那多点情况下，它的 rulebook 就会发生改变了，因为它的坐标发生了改变，新多了一个新坐标，那新多的坐标也会参与计算，因此最终的关系会变成图中的样子，这个大家稍微推一下就好了

那么所以说我们可以通过 rulebook 得到一系列的 atomic operation，我们把这一系列 atomic operation 放一起，可以把这个带有 sparse 的 convolution 很多没有意义的点最后给它全部都 skip 掉，就只保留我们真正想要做计算的点，这样可以将一个 sparse 的 convolution 计算变成 dense 的 matmul 计算

那么其实这个也就意味着如果说我们的 input size 变得非常大，如上图所示，那 input 的 size 越大 input 的稀疏性也越大，也就是输入中不为 0 的点就越少，比如上图中真正参与计算的只有 6 个点，那么我们只要把这 6 个点的信息给放入到 rulebook 里面去，之后让它去得到那个输入输出和 Filter 之间的对应关系，我们就可以把它变成一个密集的矩阵乘法计算，这个其实就是 Sparse Convolution 的一个计算原理

OK，我们稍微总结一下 3D Sparse Convolution 的特点，如下所示：

• input/output 的 tensor 是具有稀疏性的
• 将 input 和 output 中的所有点都进行计算的话，计算会很冗长
  • 有很多点是无效点，可以 skip 掉
• 需要通过某种方式建立起 input-weight-output 的关系
  • hash table：保存 input 和 output 中有数据的点
  • rulebook：保存 table 中的点和点的对应关系，以及参与计算的权重位置信息
• 计算时可以通过 rulebook 来选择性的计算
  • 可以将 sparse 的 conv 转换为 dense 的 matmul，从而得到加速

以上就是 spconv 的一个整体流程，下面我们看一下 BEVFusion 是怎么实现 spconv

BEVFusion 中使用的 spconv 的实现是节选自 CenterPoint 的实现，核心代码是：https://github.com/tianweiy/CenterPoint/blob/master/det3d/models/backbones/scn.py

@BACKBONES.register_module
class SpMiddleResNetFHD(nn.Module):
    def __init__(
        self, num_input_features=128, norm_cfg=None, name="SpMiddleResNetFHD", **kwargs
    ):
        super(SpMiddleResNetFHD, self).__init__()
        self.name = name

        self.dcn = None
        self.zero_init_residual = False

        if norm_cfg is None:
            norm_cfg = dict(type="BN1d", eps=1e-3, momentum=0.01)

        # input: # [1600, 1200, 41]
        self.conv_input = spconv.SparseSequential(
            SubMConv3d(num_input_features, 16, 3, bias=False, indice_key="res0"),
            build_norm_layer(norm_cfg, 16)[1],
            nn.ReLU(inplace=True)
        )

        self.conv1 = spconv.SparseSequential(        
            SparseBasicBlock(16, 16, norm_cfg=norm_cfg, indice_key="res0"),
            SparseBasicBlock(16, 16, norm_cfg=norm_cfg, indice_key="res0"),
        )

        self.conv2 = spconv.SparseSequential(
            SparseConv3d(
                16, 32, 3, 2, padding=1, bias=False
            ),  # [1600, 1200, 41] -> [800, 600, 21]
            build_norm_layer(norm_cfg, 32)[1],
            nn.ReLU(inplace=True),
            SparseBasicBlock(32, 32, norm_cfg=norm_cfg, indice_key="res1"),
            SparseBasicBlock(32, 32, norm_cfg=norm_cfg, indice_key="res1"),
        )

        self.conv3 = spconv.SparseSequential(
            SparseConv3d(
                32, 64, 3, 2, padding=1, bias=False
            ),  # [800, 600, 21] -> [400, 300, 11]
            build_norm_layer(norm_cfg, 64)[1],
            nn.ReLU(inplace=True),
            SparseBasicBlock(64, 64, norm_cfg=norm_cfg, indice_key="res2"),
            SparseBasicBlock(64, 64, norm_cfg=norm_cfg, indice_key="res2"),
        )

        self.conv4 = spconv.SparseSequential(
            SparseConv3d(
                64, 128, 3, 2, padding=[0, 1, 1], bias=False
            ),  # [400, 300, 11] -> [200, 150, 5]
            build_norm_layer(norm_cfg, 128)[1],
            nn.ReLU(inplace=True),
            SparseBasicBlock(128, 128, norm_cfg=norm_cfg, indice_key="res3"),
            SparseBasicBlock(128, 128, norm_cfg=norm_cfg, indice_key="res3"),
        )


        self.extra_conv = spconv.SparseSequential(
            SparseConv3d(
                128, 128, (3, 1, 1), (2, 1, 1), bias=False
            ),  # [200, 150, 5] -> [200, 150, 2]
            build_norm_layer(norm_cfg, 128)[1],
            nn.ReLU(),
        )

    def forward(self, voxel_features, coors, batch_size, input_shape):

        # input: # [41, 1600, 1408]
        sparse_shape = np.array(input_shape[::-1]) + [1, 0, 0]

        coors = coors.int()
        ret = spconv.SparseConvTensor(voxel_features, coors, sparse_shape, batch_size)

        x = self.conv_input(ret)

        x_conv1 = self.conv1(x)
        x_conv2 = self.conv2(x_conv1)
        x_conv3 = self.conv3(x_conv2)
        x_conv4 = self.conv4(x_conv3)

        ret = self.extra_conv(x_conv4)

        ret = ret.dense()

        N, C, D, H, W = ret.shape
        ret = ret.view(N, C * D, H, W)

        multi_scale_voxel_features = {
            'conv1': x_conv1,
            'conv2': x_conv2,
            'conv3': x_conv3,
            'conv4': x_conv4,
        }

        return ret, multi_scale_voxel_features

以上就是 SCN 网络的 forward 代码，前向过程无非是几个 conv 的叠加，从上面的代码中也能看到 conv1、conv2、conv3、conv4 的实现是通过 SparseSequential 这一个序列将多个块堆叠，其中包含 SubMConv3d、SparseConv3d、SparseBasicBlock 等等

值得注意的是 SCN 网络前向传播的输入尺寸是 41x1600x1200，输出尺寸是 2x200x150；而在 CUDA-BEVFusion 中它在 SCN 网络的 forward 最后又添加了 scatter 和 reshape 操作，所以 CUDA-BEVFusion 中 SCN 网络前向传播的输入尺寸是 41x1440x1440，输出尺寸是 1x256x180x180

下面我们再来看下核心的 SparseBasicBlock 的具体实现，代码如下：

class SparseBasicBlock(spconv.SparseModule):
    expansion = 1

    def __init__(
        self,
        inplanes,
        planes,
        stride=1,
        norm_cfg=None,
        downsample=None,
        indice_key=None,
    ):
        super(SparseBasicBlock, self).__init__()

        if norm_cfg is None:
            norm_cfg = dict(type="BN1d", eps=1e-3, momentum=0.01)

        bias = norm_cfg is not None

        self.conv1 = conv3x3(inplanes, planes, stride, indice_key=indice_key, bias=bias)
        self.bn1 = build_norm_layer(norm_cfg, planes)[1]
        self.relu = nn.ReLU()
        self.conv2 = conv3x3(planes, planes, indice_key=indice_key, bias=bias)
        self.bn2 = build_norm_layer(norm_cfg, planes)[1]
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = replace_feature(out, self.bn1(out.features))
        out = replace_feature(out, self.relu(out.features))

        out = self.conv2(out)
        out = replace_feature(out, self.bn2(out.features))

        if self.downsample is not None:
            identity = self.downsample(x)

        out = replace_feature(out, out.features + identity.features)
        out = replace_feature(out, self.relu(out.features))

        return out

我们可以在 SparseBasicBlock 模块的 forward 前向传播中看到它是 conv1、conv2、bn、relu 等模块的堆叠，核心就是 conv1 和 conv2，而 conv1、conv2 的实现又来自 conv3x3，conv3x3 实现代码如下所示：

def conv3x3(in_planes, out_planes, stride=1, indice_key=None, bias=True):
    """3x3 convolution with padding"""
    return spconv.SubMConv3d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=1,
        bias=bias,
        indice_key=indice_key,
    )

可以看到它最终会调用 spconv 中的 SubMConv3d 这个模块，这个模块其实就是我们前面一直在讲的稀疏卷积，值得注意的是在 spconv 中有 SparseConv3d 和 SubMConv3d 两种稀疏卷积形式：

一种是 Spatially Sparse Convolution，在 spconv 中为 SparseConv3D。就像普通的卷积一样，只要卷积核 kernel 覆盖了一个非零输入点，就会计算出对应的输出。对应论文：SECOND:Sparsely Embedded Convolutional Detection

另外一种是 Submanifold Sparse Convolution，在 spconv 中为 SubMConv3D。只有当卷积核 kernel 中心覆盖了一个非零输入点时，卷积输出才会被计算。对应论文：3D Sematic Segmentation with Submanifold Sparse Convolutional Networks

两种稀疏卷积的区别如下图所示：

从图中可以看到输入 Input 上只有 P1 和 P2 位置上是有值的，其它位置全为 0，有两一个卷积核参与计算，输出中 A1 表示是由 P1 计算出来的结果，A2 表示是由 P2 计算出来的结果，A1A2 表示是由 P1 和 P2 共同计算出来的结果

可以看到 SparseConv3D 的输出中只有一个地方没有值，因为只要卷积核覆盖一个非零输入的就会去计算，而相反 SubMConv3D 的输出中仅仅只有两个地方有值，这是因为 SubMConv3D 只有当卷积核中心点覆盖到非零输入时才会去计算，那么以上就是这两种稀疏卷积的一个区别了，相对而言还是比较好理解的

2. 导出带有spconv网络的onnx需要考虑的事情

在 BEVFusion 部署中，我们需要将 SCN 导出为 ONNX，使用 CUDA、TensorRT 进行加速部署，但值得注意的是我们并不能按照以往的方式进行 onnx 的 eport 导出，原因是 spconv 的节点并不是常规的像 conv、bn、relu 一样，这个节点它在 onnx 中并不存在，需要我们自己创建。

我们自己创建好 spconv 节点后导出也依旧存在问题，这是因为 pytorch 转 onnx 的过程中会对网络的 forward 前向传播过程进行 trace 跟踪，从而得到 onnx 相对应的计算节点，但是 spconv 内部处理复杂并且 layer 间的 tensor 的形式比较特殊，导致在 forward 过程中没有办法 trace 到内部的一些计算方法和逻辑，这也就意味着 pytorch 官方实现的 jit trace 已无法满足我们的需求，所以没有办法正常导出。

解决方案杜老师在 复杂onnx解决方案（以sparseconv为例） 课程中也提到过，自己来实现 trace，利用 python 最核心的特性，直接替换特定函数的实现，以实现挂钩到自己函数中。具体来说对 spconv.conv.SparseConvolution.forward 的实现进行重定位（hook 钩子函数实现），取出 spconv 推理需要的信息，通过 onnx_helper 创建自定义 node 和 graph，实现导出。

导出后的 onnx 模型如下图所示：

上图中的 SparseConvolution 这个节点是 onnx 中不存在的，这是自己创建的一个 onnx 节点，我们可以看到这个 onnx 节点中包含了非常多的属性，比如 kernel_size、padding、rulebook、stride 等等，我们把这些信息给它导出出去，之后在读取 onnx 的时候再把这些信息给 parse 出来，再做个前向推理，整个过程就完成了

NVIDIA-AI-IOT 中开源了 SCN 的 onnx 导出方法，里面使用到了 onnx_helper 中的 make_node、make_tensor、make_graph 等内容，具体导出实现代码在：https://github.com/NVIDIA-AI-IOT/Lidar_AI_Solution/blob/master/libraries/3DSparseConvolution/tool/centerpoint-export/exptool.py

下节课我们会详细分析导出 onnx 的代码，大家感兴趣的可以先看一下这部分

总结

这节课程我们学习了 spconv 的原理，它通过 hash table 将输入输出中不为 0 的点的位置坐标保存下来，并通过 rulebook 记录输入输出和 filter 之间的对应关系，把没有意义的点全部 skip 掉，只保留真正想做计算的点，从而将一个稀疏的卷积计算变成密集的矩阵乘法计算，实现加速。最后我们还简单的介绍了带有 spconv 的 SCN 网络导出 onnx 的难点以及对应的解决方案

OK，以上就是第 2 小节有关 spconv 原理的全部内容了，下节我们将去学习导出带有 spconv 的 SCN 网络的 onnx，敬请期待😄

下载链接

• 论文下载链接【提取码：6463】
• 数据集下载链接【提取码：data】
• 代码和安装包下载链接【提取码：cuda】

参考

• 复杂onnx解决方案（以sparseconv为例）
• SECOND:Sparsely Embedded Convolutional Detection
• 3D Sematic Segmentation with Submanifold Sparse Convolutional Networks
• https://github.com/tianweiy/CenterPoint/blob/master/det3d/models/backbones/scn.py
• https://github.com/NVIDIA-AI-IOT/Lidar_AI_Solution/blob/master/libraries/3DSparseConvolution/tool/centerpoint-export/exptool.py