【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代码解读 - 知乎