JSparse：Jittor稀疏卷积计算库开源啦！

分享
1

10月20日，计图团队发布了稀疏卷积计算库，可支持稀疏卷积，稀疏矩阵乘法等，该类能够支持3D物体分类，分割，识别和重建等任务。

计图(Jittor)是清华大学计算机系图形学实验室于2020年3月20日发布并开源的深度学习框架。计图团队已经陆续发布了GAN模型库、语义分割模型库、检测与实例分割库、三维点云库、可微渲染库、遥感检测库、NeRF模型库等。

Part 1
背景介绍

近年来，稀疏卷积方法在计算机图形学各个领域都比较活跃，常用于各类3D点云项目，如室内/外场景点云的实例分割、自动驾驶中的物体识别、对残缺点云的补全、三维物体的分类或重建工作。

image
image700×457 69.1 KB

(a) 三维物体语义分割
image
(b) 重建[1]
image
image932×491 160 KB

(c) 物体识别[2]

图1 基于体素卷积方法的应用

三维空间相较于二维图片而言，随规模提升所产生的数据量飚升，使用标准的三维卷积操作耗时较大。与此同时，由于三维点云数据在空间中分布的稀疏性，稀疏卷积应运而生，在近些年的各类任务上均取得了相当不错的效果。

Part 2

JSparse计算库介绍

稀疏卷积算法的算法原理是：先将卷积核每个偏移对应的体素位置映射成相对应的键值(gather)，然后在卷积的时候利用哈希表查询对应位置的所包含的点，并且导出该偏移的权值矩阵，对二者相乘(matmul)，最后再对求得键值进行空间上的还原(scatter)。

image
image1023×554 223 KB

图2 基于离散化方法的稀疏卷积算法[4]

Jittor团队及开源社区开发者借鉴TorchSparse库[2-3]开发了JSparse库，实现了更多的算子如池化操作，以适应多种类的卷积神经网络(例如点云分类、语义分割)。其次，开发团队对于其中的实现进行了优化，例如对quantize离散化采用速度更快的算子。最后，基于jittor高效的元算子融合技术，获得了整体上的速度提升。

如表1和表2所示，在ScanNet[4]点云数据集上进行测试下取得了10%~30%的综合效率提升，在VMNet[5]上测试约15%的训练速度提升(细节可参考仓库实现)。

表1稀疏卷积的性能对比

image
image752×298 61.5 KB

表2 VMNet的速度比较
image

JSparse支持以下算法：

  1. 稀疏卷积以及构建稀疏卷积网络所需的归一化层，激活层，池化层等模块。

  2. 稀疏矩阵乘法，通用三对角求解算法以及二元稀疏算子。

  3. 常用的点云处理算法，如FPS, KNN, Radius等。

以下代码展示了如何使用JSparse构建一个简单的稀疏卷积网络。

import jsparse.nn as spnn
from jittor import nn
algorithm = "cuda"
model = nn.Sequential(
    spnn.Conv3d(3, 32, 2),
    spnn.BatchNorm(32),
    spnn.ReLU(),
    spnn.Conv3d(32, 64, 3, stride=1, algorithm=algorithm),
    spnn.BatchNorm(64),
    spnn.ReLU(),
    spnn.Conv3d(64, 128, 3, stride=1, algorithm=algorithm),
    spnn.BatchNorm(128),
    spnn.ReLU(),
    spnn.Conv3d(128, 256, 2, stride=2, algorithm=algorithm),
    spnn.BatchNorm(256),
    spnn.ReLU(),
    spnn.Conv3d(256, 128, 2, stride=2, transposed=True, algorithm=algorithm),
    spnn.BatchNorm(128),
    spnn.ReLU(),
    spnn.Conv3d(128, 64, 3, stride=1, transposed=True, algorithm=algorithm),
    spnn.BatchNorm(64),
    spnn.ReLU(),
    spnn.Conv3d(64, 32, 3, stride=1, transposed=True, algorithm=algorithm),
    spnn.BatchNorm(32),
    spnn.ReLU(),
    spnn.Conv3d(32, 3, 2),
)

JSparse全部代码已开源在Gitlink和Github：

欢迎大家使用JSparse开展研究工作。

Part 3

开源贡献者

JSparse的开源生态发展离不开开发者的贡献，感谢腾讯犀牛鸟开源人才培养计划对JSparse的支持！JSparse计算库是在计图团队核心成员李相利博士生的带领下完成的。感谢以下开发者对JSparse作出的贡献！

1） 陈逸凡，华东师范大学三年级硕士生，数学专业
image
2） 张嘉鹏，清华大学三年级本科生，信息与计算科学专业
image
参考文献

  1. Choy C, Lee J, Ranftl R, et al., High-dimensional convolutional networks for geometric pattern recognition, Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2020, 11227-11236.
  2. Yan Y, Mao Y, Li B, Second: Sparsely embedded convolutional detection, Sensors, 2018, 18(10), 3337.
  3. Tang H, Liu Z, Zhao S, et al., Searching efficient 3d architectures with sparse point-voxel convolution, European conference on computer vision. Springer, Cham, 2020, 685-702.
  4. Tang H, Liu Z, Li X, et al., TorchSparse: Efficient Point Cloud Inference Engine, Proceedings of Machine Learning and Systems, 2022, 4, 302-315.
  5. Dai A, Chang A X, Savva M, et al., Scannet: Richly-annotated 3d reconstructions of indoor scenes, Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 5828-5839.
  6. Hu Z, Bai X, Shang J, et al. Vmnet: Voxel-mesh network for geodesic-aware 3d semantic segmentation, Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021, 15488-15498.