【OpenPCDet】稀疏卷积SPConv-v1.2代码解读(1)

【3D卷积】

        以下左图展示了一个2D卷积，使用一个3x3的卷积核，在单通道图像上进行卷积，其中Padding为1，得到输出。右图为一个单通道的3D卷积,与2D卷积不同之处在于，输入图像多了一个 depth 维度，卷积核也多了一个depth维度，之前2D卷积上3x3的卷积核现在变成了3x3x3。这里的3D不是通道导致的，而是深度（多层切片，多帧视频），因此，虽然输入和卷积核和输出都是3D的，但都可以是单通道的。
 

【3D稀疏卷积】

        标准的3D卷积直接用于类似3D点云检测/分割等3D任务时，因为该类场景中输入特征的稀疏性，会带来严重的冗余计算量。所以，面对稀疏场景发展除了3D稀疏卷积。当然，2D也有稀疏卷积。

对于稀疏卷积有两种：

一种是Spatially Sparse Convolution ，在spconv中为SparseConv3d。就像普通的卷积一样，只要kernel 覆盖一个 active input site，就可以计算出output site。

另一种是Submanifold Sparse Convolution， 在spconv中为SubMConv3d。只有当kernel的中心覆盖一个 active input site时，卷积输出才会被计算。

【Second引入3D稀疏卷积】

        Second论文中，作者在VoxleNet论文的基础上作了进一步的发展。考虑到VoxleNet模型中3D卷积运算量较大，速度不佳。作者引入了稀疏3D卷积来代替，在检测速度和内存使用方面都做了优化。作和开源了3D稀疏卷积的实现：GitHub - traveller59/spconv: Spatial Sparse Convolution Library。

截止目前已更新至spconv 2.x版本。我这里作代码解读仍然时基于早期的1.2版本，对于理解思想来说，问题不大。不得不佩服作者强大的代码工程能力!

【Second网络结构中的3D稀疏卷积】

        分析OpenPCDet中Second的网络结构，3D稀疏卷积使用在3D骨干网模块：VoxelBackBone8x中。它接收MeanVFE模块的结果，经过精心设计好的3D稀疏卷积和3D稀疏子流卷积的有效组合，得到输出特征，并送入HeightCompression作深度方向的压缩，后面就是我们熟悉的2D检测网络的结构：2D骨干网络-->RPN-->分类/回归-->后处理。

 在Second原始论文中给出了BACKBONE_3D部分的示意图，但是这与OpenPCDet中Second具体的网络参数有所差异，注意区分。下图是Second中BACKBONE_3D部分的表示。

OpenPCDet中的Second我们可以参考具体的代码实现。

class VoxelBackBone8x(nn.Module):
    def __init__(self, model_cfg, input_channels, grid_size, **kwargs):
        super().__init__()
        self.model_cfg = model_cfg
        norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
 
        self.sparse_shape = grid_size[::-1] + [1, 0, 0] #e.g. array([  41, 1600, 1408])
 
        self.conv_input = spconv.SparseSequential(
            spconv.SubMConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),
            norm_fn(16),
            nn.ReLU(),
        )
        block = post_act_block
 
        self.conv1 = spconv.SparseSequential(
            block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'),
        )
 
        self.conv2 = spconv.SparseSequential(
            # [1600, 1408, 41] -> [800, 704, 21]
            block(16, 32, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
            block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
            block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
        )
 
        self.conv3 = spconv.SparseSequential(
            # [800, 704, 21] -> [400, 352, 11]
            block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
        )
 
        self.conv4 = spconv.SparseSequential(
            # [400, 352, 11] -> [200, 176, 5]
            block(64, 64, 3, norm_fn=norm_fn, stride=2, padding=(0, 1, 1), indice_key='spconv4', conv_type='spconv'),
            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
        )
 
        last_pad = 0
        last_pad = self.model_cfg.get('last_pad', last_pad)
        self.conv_out = spconv.SparseSequential(
            # [200, 150, 5] -> [200, 150, 2]
            spconv.SparseConv3d(64, 128, (3, 1, 1), stride=(2, 1, 1), padding=last_pad,
                                bias=False, indice_key='spconv_down2'),
            norm_fn(128),
            nn.ReLU(),
        )
        self.num_point_features = 128
        ....
 
    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size: int
                vfe_features: (num_voxels, C)
                voxel_coords: (num_voxels, 4), [batch_idx, z_idx, y_idx, x_idx]
        Returns:
            batch_dict:
                encoded_spconv_tensor: sparse tensor
        """
        voxel_features, voxel_coords = batch_dict['voxel_features'], batch_dict['voxel_coords']
        pdb.set_trace()
        batch_size = batch_dict['batch_size']
        input_sp_tensor = spconv.SparseConvTensor(
            features=voxel_features,                     #e.g. torch.Size([16000, 4])
            indices=voxel_coords.int(),                  #e.g. torch.Size([16000, 4]) 
            spatial_shape=self.sparse_shape,             #e.g. array([41, 1600, 1408])
            batch_size=batch_size
        )
        x = self.conv_input(input_sp_tensor)
        x_conv1 = self.conv1(x)
        x_conv2 = self.conv2(x_conv1)   #stride 2,downsample
        x_conv3 = self.conv3(x_conv2)   #stride 2,downsample
        x_conv4 = self.conv4(x_conv3)   #stride 2,downsample
 
        # for detection head
        # [200, 176, 5] -> [200, 176, 2]
        out = self.conv_out(x_conv4)
        ....

对于VoxelBackbone8x模块的前向推理(forward)部分，其输入字典中最重要的内容为voxel_features和voxel_coords。他们分别表示有效的输入特征，以及这些有效特征的空间位置。voxel_features的size为(N,4)，通常不同帧点云的有效特征的数量是不同的，N表示当前batch中总的有效输入特征的数量。可见，就VoxelBackBone8x模块来说，输入feature map其实是不固定的。具体到无论是3D标准稀疏卷积还是3D子流形卷积，其输入特征也是可变的。这会给我们做进一步做spconv的部署带来挑战。就我们常用的TensorRT推理引擎来说，它就要求输入特征是固定的。注意，这里说的固定跟TensorRT中的dynamic shape不是一回事。

【spconv模块】

        在Second中spconv是作为Pytorch的一个自定义扩展模块在使用。对于一般的操作，我们扩展Pytorch模块很容易，只用使用Python来扩展即可。只需要继承torch.nn.Module并实现其__init__,forward等方法，求导的函数是不需要设置的，会自动按照求导规则求导师。像Second代码中的HeightCompression模块就是这样一个例子。这种扩展方式即插即用，不需要编译。

class HeightCompression(nn.Module):
    def __init__(self, model_cfg, **kwargs):
        super().__init__()
        self.model_cfg = model_cfg
        self.num_bev_features = self.model_cfg.NUM_BEV_FEATURES
 
    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                encoded_spconv_tensor: sparse tensor
        Returns:
            batch_dict:
                spatial_features:
        """
        encoded_spconv_tensor = batch_dict['encoded_spconv_tensor']
        spatial_features = encoded_spconv_tensor.dense()
        N, C, D, H, W = spatial_features.shape   #e.g. torch.Size([1, 128, 2, 200, 176])
        spatial_features = spatial_features.view(N, C * D, H, W)
        batch_dict['spatial_features'] = spatial_features
        batch_dict['spatial_features_stride'] = batch_dict['encoded_spconv_tensor_stride']
        return batch_dict

但是对于像对3D稀疏卷积这样复杂的操作进行优化，实现其扩展模块，单纯靠Pytorch已实现的operator的组合已经无法做到。正如spconv的实现，它采用C++和CUDA来扩展自定义模块。在 PyTorch 中直接扩展底层C++算子主要有三种方式，native_functions.yaml、C++ extension方式、OP register方式。spconv中使用了OP register这种方式。spconv代码主要分为python部分代码和c++/cuda部分代码两部分，对于其中重要内容后文我们做详细分析。

  python部分目录

├── setup.py
├── spconv
│   ├── conv.py
│   ├── functional.py
│   ├── identity.py
│   ├── __init__.py
│   ├── modules.py
│   ├── ops.py
│   ├── pool.py
│   ├── tables.py
│   ├── test_utils.py
│   └── utils
│       ├── __init__.py
│       └── __pycache__

c++/cuda部分目录

├── include
│   ├── cuhash
│   ├── paramsgrid.h
│   ├── spconv
│   ├── tensorview
│   ├── torch_utils.h
│   └── utility
│       └── timer.h
├── src
│   ├── cuhash
│   ├── spconv
│   │   ├── all.cc
│   │   ├── CMakeLists.txt
│   │   ├── cublas_gemm.cc
│   │   ├── indice.cc
│   │   ├── indice.cu
│   │   ├── maxpool.cc
│   │   ├── maxpool.cu
│   │   ├── pillar_scatter.cu
│   │   ├── pool_ops.cc
│   │   ├── reordering.cc
│   │   ├── reordering.cu
│   │   └── spconv_ops.cc
│   └── utils
│       ├── all.cc
│       ├── CMakeLists.txt
└── third_party

【参考文献】

稀疏卷积 Sparse Convolution Net - 知乎

PyTorch扩展自定义PyThon/C++(CUDA)算子的若干方法总结 - 知乎

这可能是关于Pytorch底层算子扩展最详细的总结了！ - 知乎

PyTorch算子底层源码解读--Op Registration - 知乎

通俗易懂的解释Sparse Convolution过程 - 知乎

Spconv代码解读 - 知乎