【代码阅读】Sparse Convolution代码理解

Sparse Convolution成功用于3D目标检测的网络，例如Second，Part-A^2，PV-RCNN等，证明其有效性。相比于3D Convolution，在运算速度和显存消耗中有巨大的优势。Sparse Convolution在SECOND论文中提出，并且原文中给出了实现的方法。但在Part-$A^2$和PV-RCNN中用的是另外一个版本，并且在pytorch中非常易用。

本文通过阅读OpenPCDet和Spconv代码，来看看spconv到底是怎么实现的。OpenPCDet主要是用来看spconv如何使用。

最重要的一句话，spconv的readme中写道：“This implementation use gather-gemm-scatter framework to do sparse convolution.”看完本文，对能理解这句话了。

OpenPCDet：VoxelBackBone8x

主要是看看spconv是怎么用的，具体查看OpenPCDet中的pcdet/models/backbones_3d/spconv_backbone.py中的VoxelBackBone8x类别。

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]  # self.sparse_shape=[41, 1600, 1480]

        self.conv_input = spconv.SparseSequential(
            spconv.SubMConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),  # 重点关注
            norm_fn(16),
            nn.ReLU(),
        )
        ...
        

	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']
        batch_size = batch_dict['batch_size']
		
		# 先初始化Sparse Tensor
        input_sp_tensor = spconv.SparseConvTensor(
            features=voxel_features,  # [32000, 3], 其中32000为两帧点云中voxel的数量，3是xyz
            indices=voxel_coords.int(),  # [32000, 4]，voxel的具体坐标，4中的第一个代表是哪一帧，后面是3维中的3个坐标
            spatial_shape=self.sparse_shape,  # self.sparse_shape=[41, 1600, 1480]，由kitti决定
            batch_size=batch_size  # batch_size=2，有两帧点云
        )

        x = self.conv_input(input_sp_tensor)  # 进行submanifold convolution
        ...
    

spconv

Sparse Tensor

在进行稀疏卷积之前，要构建Sparse Tensor，先看看这个是怎么构建的，在spconv中的spconv/__init__.py中：

class SparseConvTensor(object):
    def __init__(self, features, indices, spatial_shape, batch_size, grid=None):
        """
        Args:
            grid: pre-allocated grid tensor. should be used when the volume of spatial shape
                is very large.
        """
        self.features = features  # 储存密集的feature
        self.indices = indices   # 储存每个feature对应的voxel坐标系下的坐标
        if self.indices.dtype != torch.int32:
            self.indices.int()
        self.spatial_shape = spatial_shape
        self.batch_size = batch_size  # 储存batch size
        self.indice_dict = {}  # 储存坐标之间的对应关系
        self.grid = grid
    
	...

SubMConv3d

接着上面看，要进行submanifold convolution，输入就是Sparse Tensor，从spconv中的spconv/conv.py中查看：

class SparseConvolution(SparseModule):
    def __init__(self,
                 ndim,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 bias=True,
                 subm=False,
                 output_padding=0,
                 transposed=False,
                 inverse=False,
                 indice_key=None):
        super(SparseConvolution, self).__init__()
        ...

		# 下面的weight和bias就是本层filter拥有的权重和偏置
        self.weight = Parameter(
            torch.Tensor(*kernel_size, in_channels, out_channels))
        if bias:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)
        ...

    def forward(self, input):
        ...
		# 上面为一些准备工作，1）通过kernel size，padding等计算out_spatial_shape，处理卷积核1x1的特殊情况。
		
		# input为sparse tensor，找到本layer要依赖的上一个indice
        datas = input.find_indice_pair(self.indice_key)
        
		...
			if self.indice_key is not None and datas is not None:
                outids, _, indice_pairs, indice_pair_num, _ = datas
            else:
            	# 找到本层输入和输出的indices和它们的关系。输入的indices就是input的indices的属性。
            	# 输出的indices则是outids，indice_pairs表示indices和outids的对应关系，indice_pair_num是用于指示的中间变量。
            	# incides: [32000, 4]
            	# outids: [32000, 4]，由于submconv性质，outids和incides是一样的，如果是标准的spconv，就不一样了
            	# indices_pairs: [27, 2, 32000]，27是因为3x3x3的卷积核。2是对应关系，第0位储存indices的下标，第1位储存outids中的下标。
            	# indice_pair_num： [27]，indice_pair_num[i]==sum(indices_pairs[i,0,:])==sum(indices_pairs[i,1,:])，用于卷积过程中构造矩阵。
                outids, indice_pairs, indice_pair_num = ops.get_indice_pairs(
                    indices, batch_size, spatial_shape, self.kernel_size,
                    self.stride, self.padding, self.dilation, self.output_padding, self.subm, self.transposed, grid=input.grid)
                input.indice_dict[self.indice_key] = (outids, indices, indice_pairs, indice_pair_num, spatial_shape)
		...
		
        if self.subm:
        	# 进行spconv
            out_features = Fsp.indice_subm_conv(features, self.weight,
                                              indice_pairs.to(device),
                                              indice_pair_num,
                                              outids.shape[0])
        ...

		# 加上偏置
        if self.bias is not None:
            out_features += self.bias
        
        # 构建输出的sparse tensor
        out_tensor = spconv.SparseConvTensor(out_features, outids,
                                             out_spatial_shape, batch_size)
        out_tensor.indice_dict = input.indice_dict
        out_tensor.grid = input.grid
        return out_tensor

从上面代码中，可以看到两个问题没有解决：1）输出out_tensor的indices（也就是outids）怎么来的，两者关系indice_pairs怎么来的；2）subm_conv怎么计算的。

get_indice_pairs

先看outids和indice_pairs怎么计算的，要看spconv的spconv/ops.py中：

def get_indice_pairs(indices,
             batch_size,
             spatial_shape,
             ksize=3,
             stride=1,
             padding=0,
             dilation=1,
             out_padding=0,
             subm=False,
             transpose=False,
             grid=None):
    ...
    # 上面就是根据输入的indices，kernel size等信息计算out_shape。
    
    if grid is None:
        if ndim == 2:
            get_indice_pairs_func = torch.ops.spconv.get_indice_pairs_2d
        elif ndim == 3:
        	# 这里get_indice_pairs_func就是get_indice_pairs_3d！！！！！！！！！
            get_indice_pairs_func = torch.ops.spconv.get_indice_pairs_3d
        else:
            raise NotImplementedError
        
        # 重点关注
        # indices： [32000, 4]
        # batch_size = 2
        # out_shape = [41, 1600, 1408]
        # spatial_shape = [41, 1600, 1408]
        # ksize = [3, 3, 3]
        # stride = [1, 1, 1]
        # int(subm) = 1
        # int(transpose) = 0
        return get_indice_pairs_func(indices, batch_size, out_shape, spatial_shape, ksize,
                            stride, padding, dilation, out_padding, int(subm), int(transpose))
    ...

那么接着看get_indice_pairs_3d这个函数，要在spconv的src/spconv/all.cc中：

static auto registry =
        ...
        
        .op("spconv::get_indice_pairs_3d", &spconv::getIndicePair<3>) # get_indice_pairs_3d相关
        .op("spconv::indice_conv_fp32", &spconv::indiceConv<float>)

		...

那就得接着看spconv中的include/spconv/spconv_ops.h

template <unsigned NDim>
std::vector<torch::Tensor>
getIndicePair(torch::Tensor indices, int64_t batchSize,
	std::vector<int64_t> outSpatialShape, std::vector<int64_t> spatialShape,
	std::vector<int64_t> kernelSize, std::vector<int64_t> stride,
	std::vector<int64_t> padding, std::vector<int64_t> dilation,
	std::vector<int64_t> outPadding, int64_t _subM, int64_t _transpose) {

	// 具体就不细节展开了，可以自己去看代码
	// 做的事情就是给定indices，conv方法，kernel size，确定输出的voxel哪些不为空，也就是找到outdis。由于submconv的性质，outdis其实是和indices是一样的。
	// 有了outdis和indices，还需要知道indices中每一个voxel的27个周围voxel中哪些对应着outdis中哪个。也就是indice_pairs。

}

indice_subm_conv

这个在spconv中的spconv/functional.py中：

class SubMConvFunction(Function):
    @staticmethod
    def forward(
            ctx,
            features,
            filters,
            indice_pairs,
            indice_pair_num,
            num_activate_out):
        ctx.save_for_backward(
            indice_pairs,
            indice_pair_num,
            features,
            filters)
        # 重要
        return ops.indice_conv(features, filters, indice_pairs, indice_pair_num, num_activate_out, False, True)

    @staticmethod
    def backward(ctx, grad_output):
        indice_pairs, indice_pair_num, features, filters = ctx.saved_tensors
        input_bp, filters_bp = ops.indice_conv_backward(features, filters, grad_output, indice_pairs, indice_pair_num, False, True)
        
        return input_bp, filters_bp, None, None, None
indice_subm_conv = SubMConvFunction.apply

indice_conv要去看spconv/ops.py：

def indice_conv(features,
              filters,
              indice_pairs,
              indice_pair_num,
              num_activate_out,
              inverse=False,
              subm=False):
    if filters.dtype == torch.float32:
    	# 重要
        return torch.ops.spconv.indice_conv_fp32(features, filters, indice_pairs,
                                               indice_pair_num, num_activate_out,
                                               int(inverse), int(subm))
    elif filters.dtype == torch.half:
        return torch.ops.spconv.indice_conv_half(features, filters, indice_pairs,
                                               indice_pair_num, num_activate_out,
                                               int(inverse), int(subm))
    else:
        raise NotImplementedError

indice_conv_fp32从spconv的src/spconv/all.cc中得知，要去看include/spconv/spconv_ops.h中：

template <typename T>
torch::Tensor indiceConv(torch::Tensor features, torch::Tensor filters,
                       torch::Tensor indicePairs, torch::Tensor indiceNum,
                       int64_t numActOut, int64_t _inverse, int64_t _subM) {
  ...
  
  // kenerVolume = 27
  // 也就是说，这个循环是对3x3x3的kernel中每个分开计算的
  for (int i = 0; i < kernelVolume; ++i) {
    auto nHot = indicePairNumCpu.data<int>()[i];
    if (nHot <= 0 || (subM && i == indicePairMaxOffset)) {
      continue;
    }
    // auto timer = spconv::CudaContextTimer<>();
    // 定义输入输出的buffer
    auto outputBufferBlob =
        torch::from_blob(outputBuffer.data<T>(), {nHot, numOutPlanes}, options);
    auto inputBufferBlob =
        torch::from_blob(inputBuffer.data<T>(), {nHot, numInPlanes}, options);

	...

	// 有了indices，outids，和他俩之间关系indice_pairs，可以对于kernel的某个cell，gather出来这个cell要计算的voxel
    functor::SparseGatherFunctor<tv::GPU, T, int> gatherFtor;
    gatherFtor(tv::TorchGPU(), tv::torch2tv<T>(inputBuffer),
               tv::torch2tv<const T>(features),
               tv::torch2tv<const int>(indicePairs).subview(i, inverse), nHot);
    TV_CHECK_CUDA_ERR();

	// 矩阵乘法，对kernel中某个cell计算经过该cell提取特征的输出
    torch::mm_out(outputBufferBlob, inputBufferBlob, filters[i]);
	
	...
	// 将输出scatter到outids中
    functor::SparseScatterAddFunctor<tv::GPU, T, int> scatterFtor;
    scatterFtor(tv::TorchGPU(), tv::torch2tv<T>(output),
                tv::torch2tv<const T>(outputBuffer),
                tv::torch2tv<const int>(indicePairs).subview(i, !inverse), nHot,
                true);
    TV_CHECK_CUDA_ERR();
    }

  return output;
}

从这里，就理解了Readme中说的：“This implementation use gather-gemm-scatter framework to do sparse convolution.”

GEMM

gemm为通用矩阵乘的优化，在本代码中，直接调用torch.mm_out这个函数，依赖torch已经做好的GEMM。GEMM的理解，可以看下述博客：
https://blog.csdn.net/qq_20880415/article/details/104332743
https://zhuanlan.zhihu.com/p/66958390

对GEMM的理解就是，GEMM并未改变运算量，是通过把循环分开写，通过更多的利用cache储存来代替内存访问，从而减少时间。