Stratified Transformer for 3D Point Cloud Segmentation

用于三维点云分割的分层变换器

摘要

近年来，三维点云的分割取得了巨大的进展。目前大多数的方法都集中在局部特征的聚合上，但却不能直接模拟长距离的依赖关系。在本文中，我们提出了能够捕捉长距离背景的分层变换器，并展示了强大的泛化能力和高性能。具体来说，我们首先提出了一个新颖的关键采样策略。对于每个查询点，我们以分层的方式对附近的点进行密集抽样，对远处的点进行稀疏抽样，作为其键值，这使得模型能够扩大有效感受野，并以较低的计算成本享受远距离语境。同时，为了应对不规则点排列带来的挑战，我们提出了第一层的点嵌入来聚合局部信息，这有利于收敛并提高性能。此外，我们采用上下文相对位置编码来适应性地捕捉位置信息。最后，我们引入了一个内存效率高的实现方法，以克服每个窗口中不同的点数量的问题。大量的实验证明了我们的方法在S3DIS、ScanNetv2和ShapeNetPart数据集上的有效性和优越性。

1. 引言

如今，三维点云可以被方便地收集。它们在各种应用中表现出巨大的潜力，如自动驾驶、机器人技术和增强现实。与二维图像中的常规像素不同，三维点的排列是不规则的，这就阻碍了直接采用研究良好的二维网络来处理三维数据。因此，当务之急是探索为三维点云数据量身定做的先进方法。

大量的方法[7, 15, 16, 34, 35, 41, 52, 61, 62]已经探索了三维点云的分割，并获得了不错的性能。他们中的大多数专注于聚集局部特征，但未能明确地模拟长距离的依赖关系，这已被证明是捕捉远距离背景的关键[49]。沿着另一条研究路线，Transformer[44]可以通过自我注意机制自然地收获长距离的信息。然而，只有有限的尝试[31, 62]将Transformer应用于三维点云。点变换器[62]提出了 "矢量自留 "和 "减法关系 "来聚合局部特征，但它仍然难以直接捕获长距离的背景。体素转化器[31]是为物体检测量身定做的，对体素进行自我注意，但由于体素化，它失去了准确的位置。

不同的是，我们开发了一个高效的分割网络，使用标准的多头自我注意[44]来捕捉长距离的上下文，同时保持位置信息的完整。为此，我们提出了一个简单而强大的框架，即分层变换器。

具体来说，我们首先将三维空间划分为不重叠的立方体窗口，其灵感来自Swin Transformer[26]。然而，在Swin Transformer中，不同的窗口独立工作，每个查询标记只选择其窗口内的标记作为键，因此只关注有限的局部区域。相反，我们提出了一种分层的键值采样策略。我们不是只选择同一窗口内的附近的点作为键，而是对远处的点进行稀疏的采样。这样一来，对于每个查询点，密集的附近点和稀疏的远处点都被抽样形成钥匙，实现了显著扩大的有效感受野，而产生的额外计算可以忽略不计。例如，我们在图1中直观地展示了有效感受场（ERF）[29]，以显示长距离语境建模的重要性。在图的中间，由于无法对直接的长距离依赖关系进行建模，办公桌仅仅关注局部区域，导致错误的预测。相反，通过我们提出的分层策略，办公桌能够聚合来自远处物体的语境，如床或窗帘，这有助于纠正预测。

此外，值得注意的是，不规则的点排列对设计三维变换器构成了重大挑战。在二维图像中，斑块状的令牌可以很容易地形成空间上有规律的像素。但三维点则完全不同。在我们的框架中，每个点都被视为一个标记，我们对每个点进行点嵌入，以便在第一层汇总本地信息，这有利于更快的收敛和更强的性能。此外，我们采用有效的相对位置编码来捕捉更丰富的位置信息。它可以通过与语义特征的交互作用，在上下文中动态地产生位置偏差。另外，考虑到不同窗口中的三维点的数量变化很大，并对点数量较少的窗口造成不必要的内存占用，我们引入了一种内存效率高的实现方法，以显著减少内存消耗。

总的来说，我们的贡献有三点：

- 我们提出了Stratified Transformer，以更稀疏的方式对远处的点进行采样，扩大了有效的感受野，建立了直接的长距离依赖关系，而产生的额外计算可以忽略不计。

- 为了处理不规则的点排列，我们设计了第一层点嵌入和有效的上下文位置编码，以及一个内存效率高的实现，以建立一个为三维点云分割量身定做的强大的转化器。

- 实验表明，我们的模型在广泛采用的大规模分割数据集，即S3DIS[1]、ScanNetv2[10]和ShapeNetPart[5]上取得了最先进的结果。广泛的消融研究验证了每个组件的好处。

2. 相关工作

视觉变换器。最近，视觉变换器[44]在二维图像理解中变得流行[4, 9, 12, 13, 26, 31, 38, 42, 43, 47, 48, 57, 60, 63]。ViT[13]将每个补丁视为一个标记，并直接使用Transformer编码器来提取图像分类的特征。此外，PVT[48]提出了一个分层结构，以获得用于语义分割的金字塔特征，还提出了空间还原注意，以节省内存。另外，Swin Transformer[26]使用了基于窗口的注意力，并在连续的Transformer块中提出了移位的窗口操作。9, 12, 57]的方法进一步提出了不同的设计，以纳入长距离和全局的依赖关系。Transformer在二维中已经很流行，但在点云上仍然没有得到充分的探索。受Swin Transformer的启发，我们对三维点云采用了分层结构和移位窗口操作。在此基础上，我们提出了一个分层的策略，用于采样键，以收获长距离的上下文，并提出了几个基本的设计，以应对不规则点排列带来的挑战。

点云分割。点云分割的方法可以分为两类，即基于体素和基于点的方法。基于体素的解决方案[7, 15, 16]首先将三维空间划分为规则的体素，然后对其进行稀疏的卷积。它们产生了不错的性能，但由于体素化而受到不准确的位置信息的影响。基于点的方法[2,6,8,11,14,17-24,28,30,33-35,37,39-41,45,46,50,52,54-56,58,59,61,62]直接采用点的特征和位置作为输入，从而保持位置信息的完整。沿着这一研究思路，人们设计了不同的特征聚合方式来学习高级语义特征。PointNet及其变种[34,35]使用最大池化来聚合特征。PointConv[52]和KPConv[41]试图使用MLP或离散核点来模仿连续卷积核。Point Transformer[62]使用 "向量自我注意 "算子来聚集局部特征，并使用 "减法关系 "来生成注意权重，但它存在缺乏长程上下文和测试中各种扰动时不够稳健的问题。我们的工作是基于指向性的，并与Transformer密切相关，但有一个根本的区别：我们的工作克服了有限的有效感受场问题，并充分利用Transformer来模拟长距离的语境依赖，而不仅仅是局部聚合。

3. 我们的方法

3.1. 概述

我们的模型概述如图2所示。我们的框架是基于点的，我们使用xyz坐标和rgb颜色作为输入。我们采用了编码器-解码器的结构，其中编码器由多个阶段组成，由下采样层连接。在编码器的开始，第一层的点嵌入模块被用来进行局部聚合。然后，在每个阶段都有几个Transformer块。至于解码器，编码器的特征被逐层上采样以变得更密集，其方式类似于U-Net[36]。

图2：（a）框架概述。(b) 分层变压器块的结构。采用分层结构来获得多层次的特征。输入的点云首先要经过点嵌入模块来聚集局部结构信息。在经过几个下采样层和转换块后，特征被上采样用于分割。SSA。Stratified Self-attention。移位的SSA。具有移位窗口的SSA。最好以彩色观看。

3.2. Transformer块

Transformer块由一个标准的多头自注意模块和一个前馈网络（FFN）组成。在数以万计的点作为输入的情况下，直接应用全局自我注意会产生不可接受的内存消耗，其中N是输入点的数量。

香草版。为此，我们采用了基于窗口的自我关注。三维空间首先被划分为不重叠的立方体窗口，其中的点分散在不同的窗口中。每个查询点不需要像全局自我关注那样关注所有的点，而只需要考虑同一窗口中的邻居。多头自我关注是在每个窗口独立进行的。由于不同的窗口可能包含不同数量的点，我们把表示为第t个窗口中的点的数量。形式上，给定是头的数量，是每个头的维度，是特征维度，对于第t个窗口中的输入点x，第t个窗口中的多头自我关注被表述为

其中

是由x通过三个线性层得到的，·表示向量q和k之间的点乘。attn是注意图，y是聚合特征，进一步投射到输出特征。

请注意，上述公式只显示了单个窗口的计算，不同的窗口以同样的方式独立工作。这样一来，内存复杂度就大大降低到

，其中k是分散在每个窗口的平均点数。

为了方便跨窗口通信，我们还在两个连续的Transformer块之间将窗口移动一半的大小，与[26]相似。移位窗口的图示在补充文件中给出。

分层关键取样策略。由于每个查询点只关注其自身窗口中的局部点，香草版的Transformer块即使在移窗的情况下也会受到有效接收域的限制，如图1所示。因此，它不能捕捉到远处物体的长距离上下文依赖关系，导致错误的预测。

一个简单的解决方案是扩大三维窗口的大小。然而，随着窗口大小的增加，内存也会增加。为了以较低的内存成本有效地聚集长距离的上下文，我们提出了一个分层的策略来取样键。如图3所示，我们将空间划分为不重叠的立体窗口，窗口大小为。对于每个查询点（用绿星表示），我们找到其窗口中的点，与vanilla版本相同。此外，我们通过最远点取样（fps）对输入点进行降样，并以更大的窗口尺寸找到点。最后，密键和疏键都形成了最终的密钥，即

。请注意，重复的键点只被计算一次。

分层变换器块的完整结构如图2（b）所示。按照通常的做法，我们在每个自我注意模块或前馈网络之前使用LayerNorm[3]。为了进一步补充各窗口之间的信息交互，在连续的Transformer块中，原始窗口被移位，而大窗口被移位。这进一步提升了表7中所列的性能。

由于关键采样的分层策略，有效接收域被明显扩大，查询特征能够有效地聚合长距离的上下文。与普通版本相比，我们只是在稀疏的远端键上产生了额外的计算，这只占最终键的10%左右。

3.3. 第一层点嵌入

在第一层，我们建立一个点嵌入模块。一个直观的选择是使用线性层或MLP将输入特征投射到一个高维度。然而，我们根据经验观察到，在第一层使用线性层，收敛速度相对较慢，性能较差，如图4所示。我们注意到，来自线性层或MLP的点特征仅仅包括其自身xyz位置和rgb颜色的原始信息，但它缺乏局部的几何和背景信息。因此，在第一个转化器区块中，注意力图无法捕捉到查询和只包含原始xyz和rgb信息的键之间的高层次相关性。这对模型的表示能力和概括能力产生了负面影响。

相反，我们建议在点嵌入模块中聚合每个点的局部邻居的特征。我们尝试了各种局部聚合的方法，如最大集合和平均集合，发现KPConv的表现最好，如表5所示。令人惊讶的是，正如表4的Exp.I和II以及Exp.V和VI所建议的那样，对架构的这一微小修改带来了相当大的改进。它证明了基于Transformer的网络中初始局部聚合的重要性。请注意，与整个网络相比，单个KPConv产生的额外计算量可以忽略不计（仅2%FLOPs）。

3.4. 上下文相对位置编码

与二维空间规则的像素相比，三维点处于一个更复杂的连续空间中，对利用xyz位置提出了挑战。[32]声称位置编码对于基于Transformer的三维网络是不必要的，因为xyz坐标已经被用作输入特征。然而，尽管Transformer块的输入已经包含了xyz位置，但当深入到网络中时，细粒度的位置信息可能会在高层特征中丢失。为了更好地利用位置信息，我们采用了一种基于上下文的自适应相对位置编码方案，其灵感来自于[51]。

特别是，对于第t个窗口中的点特征X，我们将xyz坐标表示为p。因此，查询和键之间的相对xyz坐标被表述为：

为了将相对坐标映射到相应的位置编码，我们维护三个可学习的查找表

，分别对应x、y和z轴。由于相对坐标是连续的浮点数，我们将

的范围统一量化为L个离散部分，并将相对坐标映射到表的索引，即：

其中是窗口大小，是量化大小，⌊-⌋表示向下取整。

我们通过查找表格来检索与索引相对应的嵌入，并将它们相加，得到位置编码：

其中

表示表t的第idx条，

是位置编码。

实际上，查询、键和值的表是不共享的。所以我们通过添加上标来区分它们。同样地，对应于查询、键和值的位置编码分别用表示。

然后，位置编码与查询和关键特征进行点乘，得到位置偏差，然后将其添加到注意力地图中。同时，我们将价值特征与其对应的位置编码相加，然后进行加权和聚合。最后，原始方程Eq.1被更新为上下文相对位置编码（cRPE）版本：

与基于MLP的位置编码相比，相对xyz坐标通过MLP直接投射到位置偏差，cRPE通过与查询和键的点积自适应生成位置偏差，从而提供语义信息。图5显示了基于MLP和cRPE的位置偏差。它揭示了这样一个事实：由基于MLP的模型产生的位置偏差在各键之间是相似的。所以它对注意力权重的影响不大。但是对于cRPE来说，不同钥匙的位置偏差有很大的不同。此外，表4的Exp.III和IV以及Exp. 表4的Exp.V和VIII也显示了cRPE的优越性。

3.5. 下采样层和上采样层

下采样层如图6（a）所示。首先，xyz坐标经过采样和分组模块，我们首先通过最远点采样（fps）[35]对中心点进行采样，然后使用kNN查询原始点，得到分组索引。中心点的数量为原始点的1/4。同时，点特征被送入Pre-LN[53]线性投影层。此外，我们利用最大池化技术，用分组指数对投影的特征进行汇总，得到输出特征。

对于上采样层，如图6（b）所示，解码器的特征首先由Pre-LN线性层投影。我们在当前xyz坐标和之前的之间进行插值[35]。前一阶段的编码器点特征经过Pre-LN线性层。最后，我们将它们相加，得到下一个解码器特征。

4. 记忆效率的实现

在二维Swin变换器中，实现基于窗口的关注很容易，因为每个窗口中的标记数量是固定的。然而，由于三维中不规则的点排列，每个窗口中的标记物数量变化很大。一个简单的解决方案是在每个窗口中用假的令牌填充到最大的令牌数量，然后再应用掩蔽的自我注意。但是这个解决方案浪费了很多内存和计算。

相反，我们首先预先计算所有需要进行点乘的查询和密钥对。如图7(a)所示，我们使用两个索引，分别对形状的q和k进行索引，其中N表示输入点的总数。然后，我们对和索引的条目进行点乘，得到形状的注意力图attn。之后，如图7（b）所示，我们直接对attn与查询的条目进行散点软化，其中软化函数被应用于attn中与中的索引相同的条目。此外，如图7（c）所示，我们使用来索引数值v，并将其与注意力图attn相乘。最后，我们将中具有相同索引的条目相加，并将结果保存在输出特征y中。注意，每个步骤都是由一个CUDA内核实现的。所以每个步骤中的中间变量几乎不占用内存。通过这种方式，我们达到了的内存复杂度，远远低于香草实现中使用的内存复杂度。更详细的内存复杂性分析和关于分区编码实现的讨论在补充文件中给出。我们的实现比vanilla的实现节省了57%的内存。

5. 实验

5.1. 实验环境

网络结构。主要结构如图2所示。xyz坐标和rgb颜色都被用作输入。我们将初始特征维度和头数分别设置为48和3，它们在每个下采样层都会翻倍。对于S3DIS，构建了四个阶段，块深度为[2，2，6，2]。相比之下，对于ScanNetv2，我们注意到，点的数量更大。所以我们在第一层点嵌入模块的基础上增加了一个额外的降采样层。然后，再加上后面四个块深度为[3，9，3，3]的阶段。因此，ScanNetv2总共有五个阶段。

实施细节。对于S3DIS，按照以前的工作[62]，我们用4个RTX 2080Ti GPU训练了76，500次迭代。批量大小被设置为8。按照通常的做法，首先对原始输入点进行网格采样，网格大小设置为0.04米。在训练过程中，最大输入点数被设置为80,000，如果点数达到这个数字，所有额外的点数都会被丢弃。窗口大小最初设置为0.16米，在每一个下采样层之后都会增加一倍。分层抽样策略的下样比例被设置为8。除非另有说明，否则我们使用Z轴旋转、比例、抖动和降色作为数据增强。

对于ScanNetv2，我们训练了600个历时，权重衰减和批量大小分别设置为0.1和8，网格采样的网格大小设置为0.02米。在训练过程中，最多有120，000个点云被送入网络。初始窗口大小被设置为0.1m。除随机抖动外，数据增强与S3DIS上的数据增强相同。ShapeNetPart的实现细节和数据集的描述在补充文件中给出。

5.2. 结果

我们与最近最先进的语义分割方法进行了比较。表1和表2显示了S3DIS和ScanNetv2数据集的结果。我们的方法在这两个具有挑战性的数据集上都取得了最先进的性能。在S3DIS上，我们的方法明显优于其他方法，甚至比Point Transformer[62]高1.6% mIoU。在ScanNetv2上，我们的方法的验证mIoU超过了其他方法，包括基于体素的方法，差距为2.1%mIoU。在测试集上，我们的方法取得了比MinkowskiNet[7]略高的结果。潜在的原因可能是ScanNetv2中的点是相对稀疏的。因此，对于基于体素的方法来说，体素化中准确位置的损失是可以忽略不计的。但是在点比较密集的S3DIS上，我们的方法超过了MinkowskiNet，差距很大，也就是6.6%的mIoU。另外，我们的方法在验证集上比MinkowskiNet好2.1% mIoU，并且在测试中遇到各种扰动时比MinkowskiNet更稳健，如表9所示。值得注意的是，在ScanNetv2上，与基于体素的方法相比，基于点的方法首次取得了更高的性能。

另外，在表3中，为了显示泛化能力，我们还对ShapeNetPart[5]的部分分割任务进行了比较。我们的方法在类别MIOU和实例MIOU方面都超过了以前的方法，达到了新的水平。尽管我们的实例mIoU与Point Transformer相当，但我们的方法在类别mIoU方面比Point Transformer有很大优势。

5.3. 消融研究

我们进行了广泛的消融研究，以验证我们方法中每个组件的有效性，并在表4中显示结果。为了使我们的结论更有说服力，我们对S3DIS和ScanNetv2数据集进行了评估。从Exp.I到V，我们每次都增加一个组件。另外，从Exp.VI到VIII，我们通过从最终模型中删除每个组件来进行双重验证，即Exp.V.

分层变换器。在表4中，比较Exp.IV和V，我们注意到，使用分层策略，模型在S3DIS上提高了1.9% mIoU，在ScanNetv2上提高了1.2% mIoU。结合图1的视觉效果，我们注意到分层策略能够扩大有效感受野并提高性能。此外，我们还在补充文件中展示了设置不同降样比例（即4、8和16）时的效果。

第一层的点嵌入。我们比较了Exp.I和II，发现该模型在第一层点嵌入的情况下有很大程度的提高。另外，我们比较了Exp.VI和V，模型在S3DIS上获得2.0%的mIoU增益，在ScanNetv2上获得4.0%的mIoU增益。架构中的这个小修改带来了相当大的好处。

为了进一步探讨局部聚合在第一层点嵌入中的作用，我们在表5中比较了不同的局部聚合方式与线性投影。很明显，所有列出的局部聚合方法在第一层点嵌入中都比线性投影好。

上下文相对位置编码。从Exp.III到IV，使用cRPE后，在S3DIS上性能提高了2.9% mIoU，在ScanNetv2上提高了1.9% mIoU。此外，当也使用分层变压器时，通过Exp.VIII和V的比较，该模型在S3DIS上仍有4.0%的mIoU增益，在配备cRPE的ScanNetv2上有2.3%的增益。

此外，我们证明了应用cRPE对每个查询、关键或价值特征的贡献。表6显示，在任何一个特征中应用cRPE都能带来改进。当同时在查询、键和值上应用cRPE时，该模型取得了最佳性能。

此外，我们将我们的方法与第3.4节中提到的基于MLP的方法进行比较。如表6所示，我们发现基于MLP的方法（第一列）实际上与没有任何位置编码的模型（第二列）没有区别。结合图5的可视化，我们得出结论，纯粹基于xyz坐标的相对位置信息是没有帮助的，因为网络的输入点特征已经包含了xyz坐标。相反，cRPE同时基于xyz坐标和上下文特征。

移位窗口。采用移位窗口是为了补充各窗口之间的信息交互。在表7中，我们比较了我们的vanilla版本和S3DIS上的Stratified Transformer的有/无移位窗口的模型。显而易见，移位窗口在我们的框架中是有效的。此外，即使没有移动窗口，Stratified Transformer仍然产生了更高的性能，即70.1%的mIoU，与vanilla版本相比。另外，在原始窗口和大窗口上的变换也是有益的。

数据增强。数据增强在训练基于Transformer的网络中起着重要作用。在我们的框架中也是如此，如Exp.V和VII，以及Exp.II和III所示。我们还在表8中调查了每种增强的贡献。

5.4. 稳健性研究

为了显示我们模型的抗干扰能力，我们通过在测试中应用各种扰动来衡量其稳健性。按照文献[54]，我们从置换、旋转、移动、尺度和抖动等方面进行评估。如表9所示，我们的方法对各种扰动非常稳健，而以前的方法在这些情况下会出现剧烈波动。值得注意的是，我们的方法在Z轴旋转90◦时的表现甚至更好（+0.63% mIoU）。

尽管Point Transformer也采用了自我关注机制，但它产生的稳健性有限。一个潜在的原因可能是Point Transformer使用了特殊的算子设计，如 "矢量自留 "和 "减法关系"，而不是标准的多头自留。

5.5. 视觉比较

在图8中，我们视觉比较了Point Transformer、基线模型和我们的模型。它清楚地显示了我们方法的优越性。由于对长距离背景的认识，我们的方法能够识别出用黄框标出的物体，而其他方法则失败。

6. 结论

我们提出了分层变换器并取得了最先进的结果。分层策略极大地扩大了有效的接受域。此外，我们还提出了第一层点嵌入和有效的上下文相对位置编码。我们的工作回答了两个问题。首先，有可能以低计算成本建立直接的长距离依赖关系，并产生更高的性能。第二，标准转化器可以应用于三维点云，具有很强的泛化能力和强大的性能。