摘要

我们提出了一种基于八叉树结构的Transformer模型——OctFormer，用于三维点云学习。OctFormer不仅可作为三维点云分割与目标检测任务中通用且高效的主干网络，还具备线性时间复杂度，能够有效扩展至大规模点云数据。将Transformer应用于点云的核心挑战在于：注意力机制固有的二次方计算复杂度，导致计算开销巨大。为应对这一问题，已有研究将点云划分为非重叠的局部窗口，并将注意力限制在每个局部窗口内。然而，各窗口内的点数差异显著，难以在GPU上实现高效并行计算。观察到注意力机制对局部窗口形状具有较强的鲁棒性，我们提出了一种新颖的八叉树注意力机制（octree attention），该机制利用八叉树结构中经过排序打乱的键（keys）对点云进行划分，从而生成包含固定点数的局部窗口，同时允许窗口形状自由变化。此外，我们进一步引入了膨胀八叉树注意力（dilated octree attention），以进一步扩大感受野。所提出的八叉树注意力机制仅需约10行代码即可在开源库中实现，当点数超过20万时，其运行速度比其他点云注意力机制快达17倍。基于该注意力机制，OctFormer具有良好的可扩展性，在一系列三维分割与检测基准测试中均取得了当前最优性能，显著优于以往基于稀疏体素（sparse-voxel）的CNN模型以及点云Transformer模型，在效率与有效性方面均表现更优。特别地，在具有挑战性的ScanNet200数据集上，OctFormer在mIoU指标上比基于稀疏体素的CNN模型高出7.3个百分点。相关代码与训练好的模型已开源，可访问 https://wang-ps.github.io/octformer 获取。