代码:
https://github.com/mit-han-lab/flatformer
论文:
https://arxiv.org/abs/2301.08739
[FlatFormer.pdf]
Flatformer是对点云检测中的 backbone3d 部分的改进工作,主要在探究怎么高效的对点云应用transformer
具体的工作如下:
一个缩写: **PCTs 即point cloud transformers**
首先作者分析了点云transformer速度慢的原因(第三章):
最简单粗暴的是全局PCTs , 即对点云使用transformer的方式就是将每个点作为一个token,然后对一个pcd中所有点做multihead attention
比如一个pcd有10万个点,那就是10万个点做multihead attention计算,会很慢,复杂度为O(N2D),N是点数,D是每个点的特征通道数。作者分析了下,当点数到32k时,在NVIDIA A6000 GPU耗时就达到了1s,可见有多慢:
然后,为了降低耗时,有人提出了局部PCTs, 即对全局一共N个点,每个点都选择它附近的K个点做multihead attention计算,复杂度就变成了O(NK2*D),N是点数,D是每个点的特征通道数,K是一个点附近点点数。然而,局部PCTs在邻居点准备方面存在显著的开销。由于点云的稀疏性和不规则性,涉及两个主要步骤:
- 找到每个点的邻居。
- 将数据从N×D格式重构为应用MHSA所需的N×K×D格式。
这些步骤很慢,需要大量时间。例如,在VoTr模型中,为Waymo上的单个场景准备邻居的步骤需要22毫秒(即总运行时间的36%),这已经比整个CenterPoint模型更慢。对于Point Transformer(PT)模型,准备邻居的成本可能占运行时间的高达70%。局部PCTs的单个层中的此开销可能超过CenterPoint模型的总运行时间。
然后又有人提出了窗口PCTs ,
SST(Swin Transformer in Point Cloud)是一种基于窗口的点云变换器,受到了Swin Transformers在各种视觉识别任务中取得的巨大成功的启发。其中,SST是代表性的工作之一。它首先将点云投影到鸟瞰空间(bird's-eye-view, BEV),然后将BEV空间分割为形状相同且不重叠的窗口,并在每个窗口内应用MHSA(multihead attention)。与Swin Transformer类似,SST使用窗口移动来实现窗口之间的信息交换。
与图像不同,点云在空间中是稀疏且非均匀分布的。因此,每个窗口中的点数不同,并且可以相差两个数量级。由于普通的MHSA核心无法有效支持可变序列长度,SST将大小相似的窗口分批处理,并在每批中将所有窗口填充到批中的最大组大小(即填充到相同的长度)。然后,它单独在每个批次中应用MHSA。在实践中,这种填充引入了Waymo上1.7倍的计算开销。更糟糕的是,将点分配到相等的窗口还引入了显著的延迟开销:在Waymo上,每个场景需要18毫秒,甚至比CenterPoint模型的总运行时间还要慢。综上所述,填充和分区的开销使得SST不太适合硬件。
最后,作者基于上面这些问题,提出了自己模型(第四章): Flatformer
作者的核心工作可以用下面这个图来展现:
FlatFormer的基本构建块是Flattened Window Attention(FWA)。
我们只要看上面这个图的右半部分就行了,左半部分是作者用来做对对比的不好的方案
如图5r所示,FWA采用基于窗口的排序来展平点云,并将其分割为等大小的组,而不是等形状的窗口。这自然解决了组大小不平衡的问题,并避免了填充和分区的开销。
FWA然后在组内应用自注意力来提取局部特征,交替排序轴以聚合来自不同方向的特征,并移动窗口以在组之间交换特征。最后,作者提供了FWA的一个实现,进一步提高了其效率并减少了开销。