3d稀疏卷积——spconv源码剖析（二）
本文基于OpenPCDet框架中CeneterPoint算法，对spconv库中稀疏卷积源码进行剖析：
首先看OpenPCDet下的pcdet/models/backbones_3d/spconv_backbone.py
from ...utils.spconv_utils import replace_feature, spconv
继续看：pcdet/utils/spconv_utils.py
try:
import spconv.pytorch as spconv
except:
import spconv as spconv
import spconv 就是在导入安装好的spconv package。 package`目录下有__init__.py 文件，在导入spconv时__init__.py中的可执行代码会被执行
from spconv import ops, utils
from spconv.conv import (SparseConv2d, SparseConv3d, SparseConvTranspose2d,
SparseConvTranspose3d, SparseInverseConv2d,
SparseInverseConv3d, SubMConv2d, SubMConv3d)
from spconv.identity import Identity
from spconv.modules import SparseModule, SparseSequential
from spconv.ops import ConvAlgo
from spconv.pool import SparseMaxPool2d, SparseMaxPool3d
from spconv.tables import AddTable, ConcatTable, JoinTable
在导入完spconv后，可以直接使用spconv.SubMConv3d,spconv.SparseConv3d，spconv.SparseConvTensor,spconv.SparseSequential等子模块
__init__.py是Python中package的标识，定义了包的属性和方法__
- __init__.py 文件的一个主要作用是将文件夹变为一个Python模块也称为包，Python 中的每个模块的包中，都有__init__.py 文件，且该文件不能删除，否则该文件夹将不再被视为模块。
- __init__.py 文件定义了包的属性和方法。其实它可以什么也不定义；可以只是一个空文件，但是必须存在。如果 __init__.py不存在，这个目录就仅仅是一个目录，而不是一个包，它就不能被导入或者包含其它的模块和嵌套包。
在__init__.py中定义了3d稀疏卷积的核心数据结构SparseConvTensor
class SparseConvTensor(object):
def __init__(self, features, indices, spatial_shape, batch_size,
grid=None):
"""
Args:
features: [num_points, num_features] feature tensor
indices: [num_points, ndim + 1] indice tensor. batch index saved in indices[:, 0]
spatial_shape: spatial shape of your sparse data
batch_size: batch size of your sparse data
grid: pre-allocated grid tensor. should be used when the volume of spatial shape
is very large.
"""
self.features = features # 有效的特征数据
self.indices = indices # 有效的voxel网格坐标
self.spatial_shape = spatial_shape # 空间形状大小
self.batch_size = batch_size
self.indice_dict = {}
self.grid = grid
SparseConvTensor，但是本身并不是一个torch tensor，只是对稀疏Tensor的一个抽象。其内部成员features,indices和spatial_shape分别表示有效的特征数据， 有效的voxel网格坐标（即voxel空间索引）以及空间形状大小。
同时在__init__.py通过torch.ops.load_library来加载libspconv.so动态库，这样可以通过src/spconv/all.c注册的python接口来调用C++/CUDA实现的函数，后续会详细介绍。
本文基于nuscenes数据集，CenterPoint 配置参数如下： tools/cfgs/nuscenes_models/cbgs_voxel0075_res3d_centerpoint.yaml，几个主要配置参数如下：
POINT_CLOUD_RANGE: [-54.0, -54.0, -5.0, 54.0, 54.0, 3.0]
VOXEL_SIZE: [0.075, 0.075, 0.2]
MAX_POINTS_PER_VOXEL: 10
看OpenPCDet中的Centerpoint BACKBONE_3D部分的代码，下面备注的参数以nuscenes数据集得出：
class VoxelResBackBone8x(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)# 固定参数eps和momentum
self.sparse_shape = grid_size[::-1] + [1, 0, 0] # array([41, 1440, 1440]) 在原始网格的高度方向上增加了一维
# SubMConv3d:只有当kernel的中心覆盖一个 active input site时，卷积输出才会被计算
# spatial_shape:[41, 1440, 1440] --> [41, 1440, 1440]
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
# spatial_shape:[41, 1440, 1440] --> [41, 1440, 1440]
self.conv1 = spconv.SparseSequential(
SparseBasicBlock(16, 16, norm_fn=norm_fn, indice_key='res1'),
SparseBasicBlock(16, 16, norm_fn=norm_fn, indice_key='res1'),
)
# SparseConv3d:就像普通的卷积一样，只要kernel 覆盖一个 active input site，就可以计算出output site
# spatial_shape:[41, 1440, 1440] --> [21, 720, 720]
self.conv2 = spconv.SparseSequential(
block(16, 32, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res2'),
SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res2'),
)
# spatial_shape:[21, 720, 720] --> [11, 360, 360]
self.conv3 = spconv.SparseSequential(
block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res3'),
SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res3'),
)
# spatial_shape:[11, 360, 360] --> [5, 180, 180]
self.conv4 = spconv.SparseSequential(
block(64, 128, 3, norm_fn=norm_fn, stride=2, padding=(0, 1, 1), indice_key='spconv4', conv_type='spconv'),
SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res4'),
SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res4'),
)
last_pad = 0
last_pad = self.model_cfg.get('last_pad', last_pad)
# spatial_shape:[5, 180, 180] --> [2, 180, 180]
self.conv_out = spconv.SparseSequential(
spconv.SparseConv3d(128, 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
self.backbone_channels = {
'x_conv1': 16,
'x_conv2': 32,
'x_conv3': 64,
'x_conv4': 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（12000，5）：Voxel特征均值， voxel_coords（12000， 4） ：Voxel坐标的索引
# 对 voxel_features 按照 coors 进行索引，coors 在之前的处理中加入例如batch这个位置，变成了四维
voxel_features, voxel_coords = batch_dict['voxel_features'], batch_dict['voxel_coords']
batch_size = batch_dict['batch_size'] # 1
# 根据voxel特征和voxel坐标以及空间形状和batch，建立稀疏tensor
input_sp_tensor = spconv.SparseConvTensor(
features=voxel_features, # torch.Size([12723, 5])
indices=voxel_coords.int(), # torch.Size([12723, 4])
spatial_shape=self.sparse_shape, # [41, 1440, 1440]
batch_size=batch_size # 1
)
# 子流线稀疏卷积+BN+Relu spatial_shape:[41, 1440, 1440]-->[41, 1440, 1440] 通道5-->16
x = self.conv_input(input_sp_tensor)
x_conv1 = self.conv1(x) # 经两次SparseBasicBlock spatial_shape:[41, 1440, 1440]-->[41, 1440, 1440] 通道16-->16
x_conv2 = self.conv2(x_conv1) # 经子流线稀疏卷积、两次SparseBasicBlock spatial_shape:[41, 1440, 1440]-->[21, 720, 720] 通道16-->32
x_conv3 = self.conv3(x_conv2) # 经子流线稀疏卷积、两次SparseBasicBlock spatial_shape:[21, 720, 720]-->[11, 360, 360] 通道32-->64
x_conv4 = self.conv4(x_conv3) # 经子流线稀疏卷积、两次SparseBasicBlock spatial_shape:[11, 360, 360]-->[5, 180, 180] 通道64-->128
# [5, 180, 180] -> [2, 180, 180] 通道128-->128
out = self.conv_out(x_conv4) # 用的巻积形式是 SparseConv3d 而不是 SubMConv3d
batch_dict.update({
'encoded_spconv_tensor': out,
'encoded_spconv_tensor_stride': 8
})
batch_dict.update({
'multi_scale_3d_features': {
'x_conv1': x_conv1,
'x_conv2': x_conv2,
'x_conv3': x_conv3,
'x_conv4': x_conv4,
}
})
batch_dict.update({
'multi_scale_3d_strides': {
'x_conv1': 1,
'x_conv2': 2,
'x_conv3': 4,
'x_conv4': 8,
}
})
return batch_dict
features和indices 的shape分别为[N,5]和[N,4]。其中N表示有效的voxel数量。indices 的4表示batch_id，z,y,x，batch_id表示batch_size的索引，从0开始。spatial_shape经过POINT_CLOUD_RANGE和VOXEL_SIZE计算，且Z轴加1后的值为[ 41, 1440, 1440]
看下面这行代码：
self.sparse_shape = grid_size[::-1] + [1, 0, 0] # array([41, 1440, 1440])
sparse_shape 的 Z轴为什么需要加1？
参考：https://github.com/open-mmlab/mmdetection3d/issues/282
SparseEncoder将在高维度上进行下采样。因此，该参数允许高度维度可以无误差地向下采样几次，并最终满足CenterPoint的实现。
类SparseSequential代码位于：spconv/modules.py，SparseSequential类负责构建稀疏卷积序列，类似于pytorch中的nn.sequential
类的继承关系：nn.Module–>SparseModule–>SparseSequential
接下来看稀疏卷积SubMConv3d和SparseConv3d的父类SparseConvolution
SparseConvolution
spconv/conv.py中的SubMConv3d和SparseConv3d都继承自SparseConvolution，SubMConv3d和SparseConv3d主要在初始化时调用，SparseConvolution 的forward负责调度执行整个稀疏卷积；
下面代码中涉及到一些具体参数以第一层卷积conv_input的输入参数标注的。
class SparseConvolution(SparseModule):
__constants__ = [
'stride', 'padding', 'dilation', 'groups', 'bias', 'subm', 'inverse',
'transposed', 'output_padding', 'fused_bn'
]
def __init__(self,
ndim, # 数据特征维度，SubMConv3d和SparseConv3d的ndim为3
in_channels, # 输入通道
out_channels, # 输出通道
kernel_size=3, # 卷积核尺寸
stride=1, # 步长
padding=0, # 填充值
dilation=1, # 空洞卷积：卷积核各个元素之间的间隔，默认卷积方式dilation=1
groups=1, # 深度可分离卷积：groups参数将 in_channel和out_channel 按照次序分别分成了一 一对应的groups组,默认为1
bias=True, # 偏置
subm=False, # 区分是标准3d稀疏卷积还是3d子流行稀疏卷积
output_padding=0, # 输出填充,默认为0
transposed=False, # 转置卷积,默认为False
inverse=False, # 反卷积,默认为False
indice_key=None, # 索引key，子流线卷积不会改变输入输出位置索引以及输出特征图空间形状，可以字典存储起来直接利用
fused_bn=False, # conv和bn融合
use_hash=False, # 分区使用cpu和gpu计算卷积，默认为False,使用gpu
algo=ops.ConvAlgo.Native): # 3种内存分配方式，值为0，1，2
super(SparseConvolution, self).__init__()
assert groups == 1
if not isinstance(kernel_size, (list, tuple)):
kernel_size = [kernel_size] * ndim
if not isinstance(stride, (list, tuple)):
stride = [stride] * ndim
if not isinstance(padding, (list, tuple)):
padding = [padding] * ndim
if not isinstance(dilation, (list, tuple)):
dilation = [dilation] * ndim
if not isinstance(output_padding, (list, tuple)):
output_padding = [output_padding] * ndim
for d, s in zip(dilation, stride): # 必须有一个为1
assert any([s == 1, d == 1]), "don't support this."
self.ndim = ndim # 3d稀疏卷积ndim为3
self.in_channels = in_channels # 5
self.out_channels = out_channels # 16
self.kernel_size = kernel_size # 3
self.conv1x1 = np.prod(kernel_size) == 1 # 计算所有元素的乘积
self.stride = stride # 1
self.padding = padding # 1
self.dilation = dilation # 空洞卷积：卷积核各个元素之间的间隔，默认卷积方式dilation=1
self.transposed = transposed # False
self.inverse = inverse # False
self.output_padding = output_padding # 0
self.groups = groups # 深度可分离卷积：groups参数将 in_channel和out_channel 按照次序分别分成了一 一对应的groups组，默认为1
self.subm = subm # False 用于区分是标准3d稀疏卷积还是3d子流行稀疏卷积
self.indice_key = indice_key # 索引key，子流线卷积不会改变输入输出位置索引以及输出特征图空间形状，可以存储起来直接利用
self.fused_bn = fused_bn # conv和bn融合
self.use_hash = use_hash # cpu版利用哈希
self.algo = algo.value # 获取枚举标签的值，0,1,2分别对应3中内存分配方式
# torch.nn.Parameter()将一个不可训练的tensor转换成可以训练的类型parameter，并将这个parameter绑定到这个module里面
# 使用这个函数的目的也是想让某些变量在学习的过程中不断的修改其值以达到最优化。
self.weight = Parameter(
torch.Tensor(*kernel_size, in_channels, out_channels))
if bias:
self.bias = Parameter(torch.Tensor(out_channels))
else:
# 将bias参数添加到模块中，通过使用给定名字可以访问该参数.
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
n = self.in_channels
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = _calculate_fan_in_and_fan_out_hwio(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
def forward(self, input):
assert isinstance(input, spconv.SparseConvTensor)
features = input.features # （N,5）
device = features.device # cuda:0
indices = input.indices # (N, 4) 网格坐标(batch_id,z,y,x)
spatial_shape = input.spatial_shape # [41, 1440, 1440]
batch_size = input.batch_size # 4
# 1.计算输出空间形状
if not self.subm: # 普通3d稀疏卷积
if self.transposed: # False
out_spatial_shape = ops.get_deconv_output_size(
spatial_shape, self.kernel_size, self.stride, self.padding,self.dilation, self.output_padding)
else:
# 子流线稀疏卷积计算输出卷积的形状 [41, 1440, 1440]-->[21, 720, 720]
out_spatial_shape = ops.get_conv_output_size(
spatial_shape, self.kernel_size, self.stride, self.padding,self.dilation)
else: # SubMConv 卷积
# 子流线卷积输入和输出形状相同
out_spatial_shape = spatial_shape # [41, 1440, 1440]
# input.update_grid(out_spatial_shape)
# t = time.time()
if self.conv1x1: # 单独处理1x1卷积
features = torch.mm(
input.features,
self.weight.view(self.in_channels, self.out_channels))
if self.bias is not None:
features += self.bias
out_tensor = spconv.SparseConvTensor(features, input.indices,input.spatial_shape,input.batch_size)
out_tensor.indice_dict = input.indice_dict
out_tensor.grid = input.grid
return out_tensor
# input为SparseConvTensor,因为submconv3d不会改变输入输出位置索引以及输出特征图空间形状
# 如果本层的outids, indice_pairs, indice_pair_num, out_spatial_shape等与之前计算过的层相同,这里用字典存储，后面直接使用，避免重复计算
datas = input.find_indice_pair(self.indice_key)
# 2.构建Rulebook，直接使用Python调用c++接口
if self.inverse: # False 处理反卷积
assert datas is not None and self.indice_key is not None
_, outids, indice_pairs, indice_pair_num, out_spatial_shape = datas
assert indice_pair_num.shape[0] == np.prod(self.kernel_size), "inverse conv must have same kernel size as its couple conv"
else: # 非反卷积
# 如：self.indice_key = 'subm1' and datas = None
if self.indice_key is not None and datas is not None:
outids, _, indice_pairs, indice_pair_num, _ = datas
else:
# indices:[N, 4], 就是input的indices的属性，即voxel网格坐标索引
# outids: [N, 4]，由于submconv性质，outids和 incides 是一样的，如果是标准的spconv，就不一样了
# indice_pairs: [2, 27, N]，2是对应关系，2表示输入和输出两个方向，第0位储存输入indices的下标，第1位储存输出outids中的下标，27 为卷积核的volume 3x3x3，N 表示输入有效(active)特征的数量
# indice_pair_num： [27]，用于保存卷积核每一个位置上的总的计算的次数，因为是稀疏卷积卷积核上每一个元素和有效数据的运算次数可能是不同的
outids, indice_pairs, indice_pair_num = ops.get_indice_pairs(
indices, # (N, 4)
batch_size, # 4
spatial_shape, # [41, 1440, 1440]
self.kernel_size, # 3
self.stride,
self.padding,
self.dilation,
self.output_padding,
self.subm,
self.transposed,
grid=input.grid,
use_hash=self.use_hash)
# 将索引信息写入input的索引字典中
input.indice_dict[self.indice_key] = (outids,indices, indice_pairs,indice_pair_num,spatial_shape)
# 3.根据构建的Rulebook执行具体稀疏卷积计算
if self.fused_bn: # False
assert self.bias is not None
out_features = ops.fused_indice_conv(features, self.weight,self.bias,indice_pairs.to(device),indice_pair_num,
outids.shape[0], self.inverse,self.subm)
else:
if self.subm: # 子流线稀疏卷积
# Fsp.indice_subm_conv和Fsp.indice_conv经function.py中的SubMConvFunction和SparseConvFunction对象辗转还是会继续调用ops模块中的indice_conv等函数。
# 最终，他们都会以torch.ops.spconv.xx的形式调用c++扩展共享库中的api来完成任务
out_features = Fsp.indice_subm_conv(features, # 输入特征(N,5)
self.weight, # 权重(27*16*32)
indice_pairs.to(device), # [2, 27, N]
indice_pair_num, # [27]，用于保存卷积核每一个位置上的总的计算的次数
outids.shape[0], # N
self.algo # 获取枚举标签的值，0,1,2分别对应3中内存分配方式
)
else:
if self.inverse:
out_features = Fsp.indice_inverse_conv(features, self.weight, indice_pairs.to(device),
indice_pair_num, outids.shape[0], self.algo)
else:
out_features = Fsp.indice_conv(features, self.weight,indice_pairs.to(device),
indice_pair_num,outids.shape[0], self.algo)
if self.bias is not None:
out_features += self.bias
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
当非子流形卷积（普通稀疏卷积）且非转置时：
# 普通子流线卷积
"""
outputSize_i=\frac{intputSize_i+2*p_i-d_i(k_i-1)-1)}{s_i}+1
"""
def get_conv_output_size(input_size, kernel_size, stride, padding, dilation):
ndim = len(input_size)
output_size = []
for i in range(ndim):
size = (input_size[i] + 2 * padding[i] - dilation[i] * (kernel_size[i] - 1) - 1) // stride[i] + 1
if kernel_size[i] == -1:
output_size.append(1)
else:
output_size.append(size)
return output_size
out_shape由get_conv_output_size求出，输出尺寸为：
o u t p u t S i z e i = i n t p u t S i z e i + 2 ∗ p i − d i ( k i − 1 ) − 1 ) s i + 1 outputSize_i=\frac{intputSize_i+2*p_i-d_i(k_i-1)-1)}{s_i}+1 outputSizei=siintputSizei+2∗pi−di(ki−1)−1)+1
看这行代码：
datas = input.find_indice_pair(self.indice_key)
indice_key作用：
find_indice_pair函数位于spconv/__init__.py
input为SparseConvTensor,因为submconv3d不会改变输入输出位置索引以及输出特征图空间形状，如果本层的outids, indice_pairs, indice_pair_num, out_spatial_shape等与之前计算过的层相同,这里用字典存储，后面直接使用，避免重复计算
def find_indice_pair(self, key):
if key is None:
return None
if key in self.indice_dict:
return self.indice_dict[key]
return None
一般在第一次构建时，indice_key为空，只有在spconv.SparseSequential中的3个block堆叠，最后一个spconv.SubMConv3d可以复用第二个spconv.SubMConv3d的indice_key，如下列代码所示：
self.conv2 = spconv.SparseSequential(
# [41, 1408, 1600,16] <- [21, 704, 800,32]
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'),
)
SparseConvolution的forward函数输入必须是一个spconv中自定义的SparseConvTensor类型。在forward中完成稀疏卷积最重要的两个步骤：
-
构建
Rulebook，直接使用Python调用c++接口 -
根据构建的Rulebook执行具体稀疏卷积计算。使用
torch.autograd.Function进行了一层封装。Function类本身表示PyTorch的一个可导函数，只要为其定义了前向推理和反向传播的实现，就可以把它当成一个普通PyTorch函数来使用。PyTorch会自动调度该函数，合适地执行前向和反向计算。对模型部署来说，Function类有一个很好的性质：如果它定义了symbolic静态方法，该Function在执行torch.onnx.export()时就可以根据symbolic中定义的规则转换成ONNX算子。这个symbolic就是前面提到的符号函数，只是它的名称必须是symbolic而已。稀疏卷积计算由Fsp.indice_subm_conv或Fsp.indice_conv完成。Fsp.indice_subm_conv和Fsp.indice_conv经function.py中的SubMConvFunction和SparseConvFunction对象辗转还是会继续调用ops模块中的indice_conv等函数。最终，他们都会以torch.ops.spconv.xx的形式调用c++扩展共享库中的api来完成任务。
3D稀疏标准稀疏卷积和3D子流行稀疏卷积分别有SparseConv3d和SubMConv3d两个类定义。这两个类都派生自SparseConvolution。其输入参数subm用于区分是标准3d稀疏卷积还是3d子流行稀疏卷积。