TorchSparse++：面向GPU上稀疏卷积的高效训练与推理框架
分布式、并行与集群计算
2023-11-23 v1 计算机视觉与模式识别
机器学习
性能
摘要
稀疏卷积在新兴工作负载中起着关键作用，包括AR/VR中的点云处理、自动驾驶，以及推荐系统中的图理解。由于计算模式稀疏且不规则，需要专门的高性能核函数。现有GPU库为稀疏卷积提供两类数据流。gather-GEMM-scatter数据流易于实现但性能非最优，而具有重叠计算与内存访问的数据流（如implicit GEMM）性能极高但工程成本很高。本文中，我们介绍TorchSparse++，一个新的GPU库，兼具二者之长。我们创建了高效的稀疏核生成器，以低于当前最优系统十分之一工程成本生成高性能稀疏卷积核。在此基础上，我们设计了稀疏自动调优器，扩展了现有稀疏卷积库的设计空间，并为训练和推理工作负载搜索最佳数据流配置。因此，在NVIDIA A100 GPU上，TorchSparse++在推理中相较最优的MinkowskiEngine、SpConv 1.2、TorchSparse和SpConv v2分别实现2.9倍、3.3倍、2.2倍和1.7倍的实测端到端加速；在七个代表性自动驾驶基准的混合精度训练中比SpConv v2快1.2–1.3倍。它还无缝支持图卷积，相较最优图深度学习库实现2.6–7.6倍的更快推理速度。
引用
@article{arxiv.2311.12862,
title = {TorchSparse++: Efficient Training and Inference Framework for Sparse Convolution on GPUs},
author = {Haotian Tang and Shang Yang and Zhijian Liu and Ke Hong and Zhongming Yu and Xiuyu Li and Guohao Dai and Yu Wang and Song Han},
journal= {arXiv preprint arXiv:2311.12862},
year = {2023}
}
备注
MICRO 2023; Haotian Tang and Shang Yang contributed equally to this project