JSparse
Introduction
JSparse is a highperformance autodifferentiation library for sparse voxels computation and point cloud processing based on Jittor, TorchSparse and Torch Cluster.
Installation
If you use cpu version, you need to install Google Sparse Hash, and choose the convolution algorithm with "jittor"
.
The latest JSparse can be installed by
python setup.py install # or
python setup.py develop
Getting Started
Architecture
 jsparse
 nn
 functional
 modules
 utils
collate/quantize/utils.py
You can use the modules from jsparse/modules
.
Sparse Tensor
Sparse tensor (SparseTensor
) is the main data structure for point cloud, which has two data fields:

Coordinates (
indices
): a 2D integer tensor with a shape of $N \times 4$, where the first dimension denotes the batch index, and the last three dimensions correspond to quantized $x, y, z$ coordinates. 
Features (
values
): a 2D tensor with a shape of $N \times C$, where $C$ is the number of feature channels. Most existing datasets provide raw point cloud data with float coordinates. We can usesparse_quantize
(provided inJSparse.utils.quantize
) to voxelize $x, y, z$ coordinates and remove duplicates.You can also use the initialization method to automatically obtain the discretized features by turning on
quantize
optioninputs = SparseTensor(values=feats, indices=coords, voxel_size=self.voxel_size, quantize=True)
We can then use sparse_collate_fn
(provided in JSparse.utils.collate
) to assemble a batch of SparseTensor
‘s (and add the batch dimension to coords). Please refer to this example for more details.
Sparse Neural Network
We finished many common modules in jsparse.nn
such like GlobalPool
.
The neural network interface in jsparse is similar to Jittor:
import jsparse.nn as spnn
def get_conv_block(self, in_channel, out_channel, kernel_size, stride):
return nn.Sequential(
spnn.Conv3d(
in_channel,
out_channel,
kernel_size=kernel_size,
stride=stride,
),
spnn.BatchNorm(out_channel),
spnn.LeakyReLU(),
)
You can get the usage of most of the functions and modules from the example examples/MinkNet/classification_model40.py
.
BenchMark
We test several networks between JSparse(v0.5.0) and TorchSparse(v1.4.0).
Because the Jittor framework is fast, inference and training are faster than PyTorch on many operators. We also speed up quantize
with jittor’s operations and get better performance.
We test the speed on the following model and choose 10 scenes from ScanNet as the dataset.
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),
)
We finnished two versions of Sparse Convolution(completed convolution function with jittor operators or cuda).
We choose attribute batch_size = 2, total_len = 10
and run on RTX3080 to test per iteration’s speed (JSparse’s version is v0.5.0
).
JSparse  TorchSparse(v1.4.0)  

voxel_size = 0.50  20.05ms  33.66ms 
voxel_size = 0.10  25.15ms  40.40ms 
voxel_size = 0.02  81.37ms  87.42ms 
We also test the same 200 scenes of ScanNet on VMNet on JSparse and TorchSparse.
We choose attribute batch_size = 3, num_workers=16
and run on RTX Titan and Intel(R) Xeon(R) CPU E52678 v3 to test per iteration’s speed.
JSparse(cuda)  TorchSparse(v1.4.0) 

0.79s  0.92s 
If we ignore the initiation(
scannet.py
), and just test the speed of network, the speed of JSparse and TorchSparse is similar.
Acknowledgements
The implementation and idea of JSparse refers to many open source libraries, including(but not limited to) MinkowskiEngine and TorchSparse.
If you use JSparse in your research, please cite our and their works by using the following BibTeX entries:
@article{hu2020jittor,
title={Jittor: a novel deep learning framework with metaoperators and unified graph execution},
author={Hu, ShiMin and Liang, Dun and Yang, GuoYe and Yang, GuoWei and Zhou, WenYang},
journal={Science China Information Sciences},
volume={63},
number={222103},
pages={121},
year={2020}
}
@inproceedings{tang2022torchsparse,
title = {{TorchSparse: Efficient Point Cloud Inference Engine}},
author = {Tang, Haotian and Liu, Zhijian and Li, Xiuyu and Lin, Yujun and Han, Song},
booktitle = {Conference on Machine Learning and Systems (MLSys)},
year = {2022}
}
@inproceedings{tang2020searching,
title = {{Searching Efficient 3D Architectures with Sparse PointVoxel Convolution}},
author = {Tang, Haotian and Liu, Zhijian and Zhao, Shengyu and Lin, Yujun and Lin, Ji and Wang, Hanrui and Han, Song},
booktitle = {European Conference on Computer Vision (ECCV)},
year = {2020}
}