TorchSparse

TorchSparse 是一个用于点云处理的高性能神经网络库。

官网 | 论文 (MICRO 2023) | 论文 (MLSys 2022) | 演讲视频 | 文档 | PyPI 服务器

简介

点云计算已成为自动驾驶及其他应用中日益重要的工作负载。与密集的 2D 计算不同，点云卷积具有 稀疏 和 不规则 的计算模式，因此需要专用的推理系统支持以及专门的高性能内核。尽管现有的点云深度学习库已经为点云上的卷积开发了不同的数据流，但它们假设在整个模型执行过程中只使用单一的数据流。在本工作中，我们系统地分析并改进了现有的数据流。我们的最终系统 TorchSparse 在 NVIDIA A100 GPU 上分别实现了相对于最先进的 MinkowskiEngine、SpConv 1.2、TorchSparse (MLSys) 和 SpConv v2 在推理方面的 2.9 倍、3.3 倍、2.2 倍 和 1.7 倍 的端到端加速。

新闻

[2024/11] TorchSparse++ 现已通过插件支持 MMDetection3D 和 OpenPCDet! 完整示例 已提供。

[2023/11] TorchSparse++ 已被 Hao Su 教授实验室（UCSD）的 [One-2-3-45++]（论文：arXiv:2311.07885）采用，用于 3D 物体生成！

[2023/10] 我们在第 56 届 IEEE/ACM 国际微架构研讨会 (MICRO 2023) 上展示了 TorchSparse++。同时，我们也正式发布了 TorchSparse++ 的源代码。

[2023/6] TorchSparse++ 被 Hao Su 教授实验室（UCSD）的 [One-2-3-45]（论文：arXiv:2306.16928）采用，用于 3D 网格重建！

[2023/6] TorchSparse++ 已发布，并在 CVPR 2023 自动驾驶相关研讨会上进行了展示。它相比之前的最先进系统，在推理速度上提升了 1.7 至 2.9 倍。

[2023/1] Argoverse 2 数据集使用 TorchSparse 实现了其基线检测器。

[2022/8] TorchSparse 在 MLSys 2022 上进行了展示。演讲视频可在 这里 观看。

[2022/6] TorchSparse 已被 SparseNeuS 采用，用于神经表面重建。

[2022/1] TorchSparse 被 MLSys 2022 接收，其特点是自适应矩阵乘法分组和局部性感知的内存访问。

[2021/6] TorchSparse v1.4 已发布。

安装

我们提供了预编译的 torchsparse v2.1.0 包（推荐），支持不同的 PyTorch 和 CUDA 版本，以简化 Linux 系统上的安装过程。

1. 确保已安装至少 PyTorch 1.9.0：

   python -c "import torch; print(torch.__version__)"
   >>> 1.10.0
   

2. 如果您希望在 GPU 上使用 TorchSparse，请确保 PyTorch 已经安装了 CUDA：

   python -c "import torch; print(torch.version.cuda)"
   >>> 11.3
   

3. 然后可以通过运行安装脚本来找到并安装合适的 TorchSparse wheel：

   python -c "$(curl -fsSL https://raw.githubusercontent.com/mit-han-lab/torchsparse/master/install.py)"
   

如果 PyPI 服务器无法正常工作，也不用担心，您仍然可以手动下载 wheel 文件。这些 wheel 文件列在 这个网站 上。您可以利用我们的安装脚本自动确定用于索引 wheel 的版本号。例如，如果您使用 PyTorch 1.11.0 和 CUDA 11.5，版本号将为 2.1.0+torch111cu115。然后您可以根据自己的 Python 版本选择合适的 wheel。

您也可以通过以下方式从源码安装我们的库：

python setup.py install

在仓库目录下，或者使用

pip install git+https://github.com/mit-han-lab/torchsparse.git

而无需克隆整个仓库。

基准测试

推理基准测试

TorchSparse 在 3D 物体检测和 LiDAR 分割基准测试中，显著优于现有的点云推理引擎，覆盖了 Pascal、Turing 和 Ampere 三代 GPU 架构以及 FP16、TF32 和 FP32 所有精度。它比最先进的 SpConv 2.3.5 快至 1.7 倍，在云端 GPU 上比 TorchSparse-MLsys 快至 2.2 倍。此外，它还将 Orin 上 SpConv 2.3.5 的延迟降低了 1.25 倍。

训练基准测试

TorchSparse 在混合精度训练速度方面优于 MinkowskiEngine、TorchSparse-MLSys 和 SpConv 2.3.5。具体而言，它在 Tesla A100 上比最先进的 SpConv 2.3.5 快 1.16 倍，在 RTX 2080 Ti 上快 1.27 倍。它还在 A100 和 2080 Ti 上的七项基准测试中，以 4.6–4.8 倍 的优势显著超越了 MinkowskiEngine。测量时批次大小为 2。

您可以在 这个链接 找到我们的基准测试结果。如需访问预处理过的数据集，请联系作者。由于许可要求，我们无法公开 SemanticKITTI、nuScenes 和 Waymo 的原始数据。

团队

TorchSparse 由以下优秀的团队开发：

• Haotian Tang：MIT EECS 博士生（2020 年至今），项目负责人，v2.0 和 v2.1 主要负责人；
• Shang Yang：MIT EECS 博士生（2023 年至今），项目负责人，v2.1 主要负责人；
• Zhijian Liu：MIT EECS 博士生（2018 年至今），项目负责人，v2.0 主要负责人；
• Xiuyu Li：UC Berkeley EECS 博士生（2022 年至今），v2.0 主要负责人；
• Ke Hong：清华大学 EE 研究生（2021 年至今），v2.1 核心开发者，负责 PCEngine 内核的编写；
• Zhongming Yu：UCSD CS 博士生（2022 年至今），v2.1 核心开发者，负责 PCEngine 内核的编写；
• Yujun Lin：MIT EECS 博士生（2018 年至今），v2.0 核心开发者；
• Yingqi Cao：UC San Diego 本科生，目前正致力于将 TorchSparse++ 集成到算法框架中；
• Guohao Dai：上海交通大学副教授，项目的导师；
• Yu Wang：清华大学教授，项目的导师；
• Song Han：MIT EECS 副教授，项目的导师。

引用

如果您使用 TorchSparse，请使用以下 BibTeX 条目进行引用：

TorchSparse++（TorchSparse v2.1）已在 MICRO 2023 上发表：

@inproceedings{tangandyang2023torchsparse,  
  title={TorchSparse++: 面向 GPU 的稀疏卷积高效训练与推理框架},  
  author={唐浩天、杨尚、刘志坚、洪科、俞仲明、李秀宇、戴国豪、王宇、韩松},  
  booktitle={IEEE/ACM 国际微架构研讨会（MICRO）},  
  year={2023}
}

TorchSparse++ 的初步版本（TorchSparse v2.1）已在 CVPR 工作坊 2023 上发表：

@inproceedings{tangandyang2023torchsparse++,
  title = {{TorchSparse++：高效点云引擎}},
  author = {唐浩天、杨尚、刘志坚、洪科、俞仲明、李秀宇、戴国豪、王宇、韩松},
  booktitle = {计算机视觉与模式识别工作坊（CVPRW）},
  year = {2023}
}

TorchSparse 已在 MLSys 2022 上发表：

@inproceedings{tang2022torchsparse,
  title = {{TorchSparse：高效点云推理引擎}},
  author = {唐浩天、刘志坚、李秀宇、林宇俊、韩松},
  booktitle = {机器学习与系统大会（MLSys）},
  year = {2022}
}

TorchSparse 的初始版本包含于 ECCV 2020 的 SPVNAS 论文中：

@inproceedings{tang2020searching,
  title = {{利用稀疏点—体素卷积搜索高效 3D 架构}},
  author = {唐浩天、刘志坚、赵圣宇、林宇俊、林继、王涵睿、韩松},
  booktitle = {欧洲计算机视觉大会（ECCV）},
  year = {2020}
}

PCEngine 论文已被 MLSys 2023 接受：

@inproceedings{hong2023pcengine,
  title={{利用硬件资源与自适应数据流实现 3D 点云中的高效稀疏卷积}},
  author={洪科、俞仲明、戴国豪、杨欣浩、连耀秀、刘泽昊、徐宁毅、王宇},
  booktitle={第六届机器学习与系统大会（MLSys）},
  year={2023}
}

致谢

我们感谢 TuSimple 的 Yan Yan 提供的有益讨论。同时，也请参阅 dgSparse 库，该库专为图和点云上的快速高效稀疏计算而设计。PCEngine（MLSys 2023）团队的工作与我们的研究密切相关。

TorchSparse 受到众多现有开源库的启发，其中包括但不限于 MinkowskiEngine、SECOND 和 SparseConvNet。

此外，我们还要感谢 AttributeDict，它提供了一种优雅的方式来管理内核/模型配置。

TorchSparse 快速上手指南

TorchSparse 是一个专为点云处理设计的高性能神经网络库，特别适用于自动驾驶等场景。它针对点云卷积的稀疏和不规则计算模式进行了优化，在 NVIDIA GPU 上相比 MinkowskiEngine、SpConv 等主流库可实现显著的端到端加速（最高达 2.9 倍）。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux
• Python 版本：建议 Python 3.7+
• PyTorch 版本：至少 1.9.0 (torch >= 1.9.0)
• CUDA 支持：若需使用 GPU 加速，必须安装带有 CUDA 支持的 PyTorch 版本。

您可以运行以下命令检查当前环境：

# 检查 PyTorch 版本
python -c "import torch; print(torch.__version__)"

# 检查 CUDA 版本 (若返回 None 则未安装 CUDA 版 PyTorch)
python -c "import torch; print(torch.version.cuda)"

安装步骤

推荐直接使用官方提供的预编译 Wheel 包进行安装，以避免复杂的源码编译过程。

方法一：自动脚本安装（推荐）

运行以下一键安装脚本，它会自动检测您的 PyTorch 和 CUDA 版本，并从服务器下载匹配的 Wheel 包：

python -c "$(curl -fsSL https://raw.githubusercontent.com/mit-han-lab/torchsparse/master/install.py)"

注意：如果上述脚本因网络原因无法连接服务器，您可以手动访问 Wheel 包列表页。根据脚本输出的版本号规则（例如 2.1.0+torch111cu115 对应 PyTorch 1.11 + CUDA 11.5），下载对应的 .whl 文件后使用 pip install <文件名> 安装。

方法二：源码安装

如果您需要最新开发版或自定义编译，可以选择从源码安装：

# 方式 A: 克隆仓库后安装
git clone https://github.com/mit-han-lab/torchsparse.git
cd torchsparse
python setup.py install

# 方式 B: 直接通过 pip 安装 (无需克隆)
pip install git+https://github.com/mit-han-lab/torchsparse.git

基本使用

TorchSparse 的 API 设计与 PyTorch 高度兼容。以下是一个最简单的示例，展示如何创建稀疏张量并执行稀疏卷积操作。

import torch
import torchsparse
import torchsparse.nn as spnn

# 1. 准备数据
# coords: 坐标矩阵 [N, 4] (batch_idx, x, y, z)
coords = torch.IntTensor([
    [0, 0, 0, 0],
    [0, 0, 0, 1],
    [0, 0, 1, 0],
    [0, 1, 0, 0]
])
# feats: 特征矩阵 [N, C]
feats = torch.FloatTensor([
    [1.0],
    [2.0],
    [3.0],
    [4.0]
])

# 2. 创建稀疏张量
input_tensor = torchsparse.SparseTensor(feats, coords)

# 3. 定义稀疏卷积层
# 输入通道 1, 输出通道 4, 核大小 3, 步长 1
conv = spnn.Conv3d(1, 4, kernel_size=3, stride=1)

# 4. 前向传播
output_tensor = conv(input_tensor)

print(output_tensor)
print(f"输出特征形状：{output_tensor.F.shape}")
print(f"输出坐标形状：{output_tensor.C.shape}")

核心概念说明：

• SparseTensor: 核心数据结构，包含特征 F (Features) 和坐标 C (Coordinates)。
• spnn.Conv3d: 稀疏 3D 卷积层，用法与 torch.nn.Conv3d 类似，但专门针对稀疏数据优化。
• 该库已集成至 MMDetection3D 和 OpenPCDet 等主流框架，可直接在这些项目中调用。