23FPS! 超强3D目标检测体素Transformer. 基于线性散布注意力的体素序列模型

3D视觉工坊 · 公众号 · · 2024-08-14 07:00

正文

点击下方卡片，关注 「3D视觉工坊」 公众号
选择星标，干货第一时间送达

来源：3D视觉工坊

添加小助理：dddvision，备注：方向+学校/公司+昵称，拉你入群。文末附3D视觉行业细分群

扫描下方二维码，加入 3D视觉知识星球，星球内凝聚了众多3D视觉实战问题，以及各个模块的学习资料：近20门视频课程（星球成员免费学习）、最新顶会论文、计算机视觉书籍、优质3D视觉算法源码等。想要入门3D视觉、做项目、搞科研，欢迎扫码加入！

背景

随着自动驾驶和智能机器人技术的发展，准确可靠的3D目标检测变得至关重要。点云数据通常由激光雷达（LiDAR）获取，具有稀疏性和不均匀分布的特点，这使得处理和理解这些数据变得复杂。现有的基于窗口的Transformer在捕捉上下文信息方面表现良好，但由于点云的稀疏性，导致每个窗口中的体素数量存在显著差异。这种差异使得现有方法需要对窗口内的体素进行分组和填充，增加了计算和内存开销。为了解决这些问题，本文提出了一种新的架构——ScatterFormer。其核心模块是散布线性注意力（Scattered Linear Attention），旨在直接对场景中的体素序列应用注意力机制，而无需将体素组织成固定长度的序列。这种方法不仅提高了计算效率，还减少了不必要的内存分配和排列操作。ScatterFormer在公开数据集Waymo和NuScenes上均取得了SOTA（State of the Art）的效果，性能优于HEDNet和DSVT。同时，ScatterFormer在A100显卡上能够达到23FPS的运行速度，显示了其在实际应用中的潜力和优势。

基于窗口的体素Transformer

与自然图像不同，点云场景中的目标通常只占据场景的一小部分区域，因此基于窗口的体素Transformer被广泛应用。最早的相关工作可以追溯到2022年的SST、VoxSeT等一系列研究。然而，由于体素的稀疏性，点云场景无法像图像一样直接将窗口化的特征展开成结构化的张量。

SST通过将窗口内的体素分组为不同的批次，采用串行-并行处理的方式来管理注意力计算。VoxSeT则是利用一组引导潜码将窗口内长度不一的体素映射到一组定长的代理体素上。2023年的DSVT和FlatFormer则采用了更加高效的窗口体素管理机制。DSVT提出了旋转集（Rotated Set），将窗口内的体素按照固定方向进行切分，形成不同的子集，并对横纵两个方向进行交替计算。FlatFormer也是类似的思路，将窗口内的体素进行子集划分。无论是早期的SST，还是近期的DSVT和FlatFormer，其核心思想都是通过划分定长集合，使得能被多头自注意力（MHSA）并行处理。但缺点是需要对体素进行多次排序，同时对于不满足定长的子集进行填充。

ScatterFormer的出发点在于，是否可以将场景体素视为一个长序列，然后对该序列的不同体素段（segment）采用独立的注意力计算？ScatterFormer的核心模块是散布线性注意力（Scattered Linear Attention），它直接对场景中的体素序列应用注意力机制，而无需将体素组织成固定长度的序列。

散布线性注意力(Scattered Linear Attention)

散布操作大致可以描述为基于给定的“组索引”对张量进行归约操作，具体可参考 https://github.com/rusty1s/pytorch_scatter 。ScatterFormer将这一形式拓展到Attention的计算上，并借助线性注意力（Linear Attention）的Softmax-free性质，提出一种散布线性注意力（Scattered Linear Attention），用于高效处理非均匀分段序列。

线性注意力机制 ：线性注意力通过利用核函数来实现对的逼近。对于窗口化的体素，可以利用线性注意力先算出每个窗口中的矩阵，这个矩阵可以理解为包含该窗口上下文信息的一个hidden state。该矩阵的大小是，与体素的数量无关。在输出特征时，只需将每个体素对应的 query 向量与 hidden state 矩阵相乘，并进行归一化即可。对于第 j 个窗口，其体素的输出特征可以写成：

其中，是第 j 个窗口中体素数量，是核函数ReLU。

在GPU上的实现 ：受到FlashAttention的启发，Scattered Linear Attention采用分块矩阵乘法来实现。首先会将体素序列切成固定大学的矩阵块，这些矩阵块会被事先放入GPU的共享内存中。GPU会为每一个窗口分配一只线程，该线程会快速遍历属于该窗口的矩阵快，累加以获得每个窗口的hidden state matrix。对于输出特征，GPU则是为每个矩阵块分配一只线程，将query矩阵与所对于的hidden state matrix进行乘法计算。

速度和内存开销 ：作者采用Triton来实现Scattered Linear Attention。具体代码放在 https://github.com/skyhehe123/ScatterFormer/blob/master/pcdet/models/model_utils/scatterformer_utils.py 。同时作者也比较了基于分块计算和基于torch_scatter的逐点计算在不同数量体素上的延迟和内存开销。由于每个矩阵块的计算都是并行触发，所以该模块的延迟增长稳定，内存与窗口中的体素数量呈线性关系。在A100 GPU上，Scattered Linear Attention在处理十万体素量级的计算时，能够在1毫秒内完成。

跨窗口交互

与之前采用shifted window来实现跨窗口交互不同，ScatterFormer采用了类Inception架构的分组卷积来实现，用两个较长的1D正交卷积核来近似大卷积核，这样既简单又高效。同时为了弥补线性注意力的局部建模能力，其中一个支路利用卷积来提取特征。其主干网络由一系列ScatterFormer Block构成，每个模块对体素特征分别进行散布线性注意力计算，跨窗口交互和基于前馈网络的通道混合。

实验

数据集 ：我们在室外场景3D点云检测数据集对Scatterformer进行了评估，其中Waymo开放数据集包含230,000个标注样本，其中160,000个用于训练，40,000个用于验证，30,000个用于测试。该数据集使用两种3D目标检测指标：平均精度均值（mAP）和按方向精度加权的平均精度均值（mAPH），并进一步分类为检测到超过五个LiDAR点的Level 1（L1）和检测到至少一个点的Level 2（L2）。NuScenes数据集包含40,000个标注样本，其中28,000个用于训练，6,000个用于验证，6,000个用于测试。在该数据集上，模型性能通过多个距离阈值（0.5、1、2和4米）的平均精度均值（mAP）和nuScenes检测分数（NDS）来衡量。 实现细节 ：Scatterformer使用开源框架OpenPCDet搭建。在Waymo数据集，体素大小设置为(0.32米, 0.32米, 0.1875米)，在NuScenes数据集上设置为(0.3米, 0.3米, 8米)。这两个数据集的窗口大小分别设置为和。我们为骨干网络堆叠了六个Scatterformer基础模块，并将注意力模块配置为具有4个头，维度为128。ScatterFormer在Waymo数据集上训练了24个周期，学习率为0.006，在NuScenes数据集上训练了20个周期，学习率为0.004。实验在一台8卡A6000的主机上进行。