CVPR2022点云语义分割：Stratified Transformer for 3D Point Cloud Segmentation

Stratified Transformer for 3D Point Cloud Segmentation

本文中的Stratified Transformer是基于Transformer的点云语义分割模型，

该模型由多个下采样层和Transformer模块组成，其中第一层的下采样层由Point Embedding模块替换。使用Point Embedding模块的目的在于，作者认为在Point Embedding模块中聚合点的局部信息有助于第一个Transformer块捕获query和key之间的高层关系。上下采样网络如下图所示：

SSA: Stratified Self-attention

和其他Transformer模块相同，作者计算了QKV三个权重，这三个权重的尺寸都为
其中，kt为第t个window中点的个数，Nh为head的个数，Nd为每个head的尺寸。

作者认为，由于每个query只关注自己窗口中的局部点，所以普通版本的Transformer块即使在窗口移位的情况下，有效接受域也有限，这会导致错误的预测。因此，作者提出了Stratified Key-sampling策略，如下图所示。

该策略将每个query的key采样分为两个分支：
第一个分支将空间划分为大小为Swin的windows，在其中寻找key_1；
第二个分支通过FPS进行下采样，使用更大的Swin_large来划分空间，并在其中寻找key_2，并将key_1和key_2合并。该模块有效的提高了模型的感受野，使得query feature可以有效聚合long-range上下文信息。

Contextual Relative Position Encoding

尽管Transformer块的输入已经包含xyz位置，但当网络变深时，细粒度的位置信息可能会在高层特征中丢失。为了更好地利用位置信息，作者采用了一种基于上下文的自适应相对位置编码方案。首先求query和key之间的相对位置：

为了将相对坐标映射到相应的位置编码，维护了三个可学习的查找表Tx，Ty，Tz。将相对坐标ri,j,m映射为表的索引：

其中Swin是是窗口大小，Squant为量化尺寸。通过查找table来检索对应的索引嵌入，然后求和得到的位置编码：

XYZ分别对应一个table。query,key和value的table不共享，因此有三组table。

Memory-efficient Implementation

内存的高效实现方法如上图所示。