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
内存的高效实现方法如上图所示。
|