[论文解读] TorchSparse: Efficient Point Cloud Inference Engine

Haotian Tang, Zhijian Liu+3 more|arXiv (Cornell University)|Apr 21, 2022

Advanced Neural Network Applications被引用 40

一句话总结

TorchSparse 在 GPU 上通过提升计算规整性和减少数据移动，加速稀疏卷积的3D点云处理，在多个模型/数据集上实现对 MinkowskiEngine 和 SpConv 的端到端速度提升约1.6x。

ABSTRACT

Deep learning on point clouds has received increased attention thanks to its wide applications in AR/VR and autonomous driving. These applications require low latency and high accuracy to provide real-time user experience and ensure user safety. Unlike conventional dense workloads, the sparse and irregular nature of point clouds poses severe challenges to running sparse CNNs efficiently on the general-purpose hardware. Furthermore, existing sparse acceleration techniques for 2D images do not translate to 3D point clouds. In this paper, we introduce TorchSparse, a high-performance point cloud inference engine that accelerates the sparse convolution computation on GPUs. TorchSparse directly optimizes the two bottlenecks of sparse convolution: irregular computation and data movement. It applies adaptive matrix multiplication grouping to trade computation for better regularity, achieving 1.4-1.5x speedup for matrix multiplication. It also optimizes the data movement by adopting vectorized, quantized and fused locality-aware memory access, reducing the memory movement cost by 2.7x. Evaluated on seven representative models across three benchmark datasets, TorchSparse achieves 1.6x and 1.5x measured end-to-end speedup over the state-of-the-art MinkowskiEngine and SpConv, respectively.

研究动机与目标

• 推动在增强现实/虚拟现实和自动驾驶场景中对3D点云进行高效实时推理。
• 解决GPU上稀疏卷积中的不规则计算和数据移动瓶颈。
• 提出一个针对自适应批处理和局部性感知数据访问进行优化的系统以提升性能。

提出的方法

• 引入自适应矩阵乘法分组，以以FLOPs换取规整性并提升GPU利用率。
• 应用量化、向量化和融合内存访问，通过优化scatter/gather来减少数据移动。
• 融合映射操作并使用局部性感知的排序以在推理期间最大化数据重用。
• 在实现CUDA后端以实现快速稀疏卷积推理的同时，提供类似PyTorch的API。

实验结果

研究问题

• RQ1 自适应矩阵乘法分批是否能提高GPU上稀疏卷积的利用率和速度？
• RQ2 数据移动优化（向量化/散点-聚集、本地性感知访问）在多大程度上降低稀疏卷积的运行时？
• RQ3 在常见的3D点云基准测试中，TorchSparse 与最先进的稀疏引擎（MinkowskiEngine、SpConv）相比如何？
• RQ4 该方法在不同数据集（如SemanticKITTI、nuScenes、Waymo）和模型（如MinkUNet、CenterPoint）上是否具有鲁棒性？

主要发现

• TorchSparse 在评估的模型/数据集上对 MinkowskiEngine 实现1.6x的端到端加速，对 SpConv 实现1.5x。
• 对矩阵乘法工作负载进行自适应分组，在矩阵乘法中实现1.4-1.5x的加速，在不同设置下总体提升为1.6x至2.0x。
• 数据移动优化（向量化 scatter/gather、带局部性感知内存访问的FP16量化、融合内核）带来显著的DRAM访问减少和加速（数据移动最高约1.9x）。
• 映射与局部性感知输入/输出访问的融合显著降低映射开销，从而在检测器中带来总体提升（最高约2.3x）。
• TorchSparse 在多GPU上针对若干MinkUNet/CenterPoint配置实现实时推理（≥10 FPS），包括在nuScenes和Waymo数据集上的3帧模型。
• 用于分批的自适应数据集/硬件特定调优（ε,S）将矩阵乘法性能提升最多约1.5x，并在各种配置下提升整体吞吐量。