3D稀疏卷积理解与使用

通俗易懂的解释Sparse Convolution过程
 sparse conv稀疏卷积
 3d稀疏卷积——spconv源码剖析系列（1-6）
3D稀疏卷积粗略理解
 spconv官方使用技巧描述
 spconv官方源代码

常用的稀疏卷积分类

Layer APIs	Common Usage	Dense Version	Note
`spconv.SparseConv3d`	Downsample	`nn.Conv3d`	Use `indice_key` to save data for inverse
`spconv.SubMConv3d`	Convolution	N/A	Use `indice_key` to save data for reuse
`spconv.SparseInverseConv3d`	Upsample	N/A	Use pre-saved `indice_key` to upsample
`spconv.SparseConvTranspose3d`	Upsample (for generative model)	`nn.ConvTranspose3d`	VERY SLOW and CAN’T RECOVER ORIGIN POINT CLOUD
`spconv.SparseMaxPool3d`	Downsample	`nn.MaxPool3d`	Use `indice_key` to save data for inverse
`spconv.SparseSequential`	Container	`nn.Sequential`	support layers above and `nn.ReLU, nn.BatchNorm, ...`
`spconv.SparseGlobalMaxPool`	global pool	N/A	return dense tensor instead of SparseConvTensor
`spconv.SparseGlobalAvgPool`	global pool	N/A	return dense tensor instead of SparseConvTensor

如何使用

初始化一个SparseConvTensor

特征：[N, num_channels]
索引：[N, (batch_idx + x + y + z)]带有批处理轴的坐标张量，坐标 xyz 顺序必须匹配空间形状和卷积参数

import spconv.pytorch as spconv
features = # 储存密集的feature [N, num_channels]
indices = # 储存每个feature对应的voxel坐标系下的坐标 [N, ndim + 1], batch 索引必须放在indices[:, 0]
spatial_shape = # 稀疏张量的空间形状, spatial_shape[i] 是indices[:, 1 + i]的形状.如spatial_shape=[60,1000,1000]代表z,y,x的空间形状，则indices[:,1:4]代表z,y,x坐标信息
batch_size = # 稀疏张量的batch_size.
x = spconv.SparseConvTensor(features, indices, spatial_shape, batch_size)
x_dense_NCHW = x.dense() # 转换稀疏张量成dense NCHW tensor.

features, indices, spatial_shape, batch_size传入之后都是可以在外部进行调用的，比如有SparseConvTensor对象sinput，使用sinput.dense(),sinput.features,sinput.indices等即可取出数据。

使用两种3D稀疏卷积

常用的两种3D稀疏卷积SparseConv3d(稀疏卷积)和SubMConv3d(子流形卷积)：

普通的卷积一样，只要kernel 覆盖一个 active input site，就可以计算出output site，对应论文SECOND: Sparsely Embedded Convolutional Detection
只有当kernel的中心覆盖一个 active input site时，卷积输出才会被计算。对应论文3D Semantic Segmentation with Submanifold Sparse Convolutional Networks

SubMConv3d输入与输出feature map上不为空的位置相同，保持了稀疏性(sparity=不为空的位置/所有位置和），也就保持了计算量。而SparseConv3d会增加稀疏性，从而也就增加了计算量。但是如果只用SubMConv3d，卷积核的感受野会限制在一定范围内，所以要结合stride=2的SparseConv3d一起使用，在尽量保持稀疏性的同时增大感受野。

class SparseConv3d(SparseConvolution):
    def __init__(self,
                 in_channels, 
                 out_channels, 
                 kernel_size, 
                 stride=1, 
                 padding=0, 
                 dilation=1, 
                 groups=1, 
                 bias=True, 
                 indice_key=None, 
                 use_hash=False, 
                 algo=ops.ConvAlgo.Native):

class SubMConv3d(SparseConvolution):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 bias=True,
                 indice_key=None,
                 algo: Optional[ConvAlgo] = None,
                 fp32_accum: Optional[bool] = None,
                 name=None):

class SparseSequential(SparseModule):
    def __init__(self, *args, **kwargs):
        super(SparseSequential, self).__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            for key, module in args[0].items():
                self.add_module(key, module)
        else:
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)
        for name, module in kwargs.items():
            if sys.version_info < (3, 6):
                raise ValueError("kwargs only supported in py36+")
            if name in self._modules:
                raise ValueError("name exists.")
            self.add_module(name, module)
        self._sparity_dict = {}

使用spconv.SparseSequential构建3D稀疏卷积网络更便捷

类似于nn.Sequential,都是将多个模块连接起来，将上一个模块的输出作为输入传入下一个模块。且SparseSequential是可以传入torch.nn中的模块的，内部做了封装，包括取出feature，feature赋值replace_feature等等，直接将SparseConvTensor给传入torch.nn的模块是不行的
，用于实现不同的操作，使用方法如下面三种:

# Example of using Sequential
model = SparseSequential(
            SparseConv2d(1,20,5),
            nn.ReLU(),
            SparseConv2d(20,64,5),
            nn.ReLU()
        )

# Example of using Sequential with OrderedDict
model = SparseSequential(OrderedDict([
            ('conv1', SparseConv2d(1,20,5)),
            ('relu1', nn.ReLU()),
            ('conv2', SparseConv2d(20,64,5)),
            ('relu2', nn.ReLU())
        ]))

# Example of using Sequential with kwargs(python 3.6+)
model = SparseSequential(
            conv1=SparseConv2d(1,20,5),
            relu1=nn.ReLU(),
            conv2=SparseConv2d(20,64,5),
            relu2=nn.ReLU()
        )

replace_feature方法

实际上就是将一个tensor读入，替换原来的feature，比如

对于F.relu，我们要传入features，之前要这样写

x.features = F.relu(x.features)

但这样写会导致torch.fx出问题，replace_feature就是为了解决这个问题而出现的

x = x.replace_feature(F.relu(x.features))

注意该函数并非在x上进行更改了，而是创建了一个全新的SparseConvTensor，也就是说要用x去接收返回值

下面给一个示例：

def forward(self, x):
    identity = self.layers_in(x)
    output = self.layers(x)
    return output.replace_feature(F.leaky_relu(output.features + identity.features, 0.1))