基于局部邻域Transformer的点云特征提取方法

张海博, 沈 洋,, 许 浩, 包艳霞,3, 刘 江

(1.浙江理工大学计算机科学与技术学院, 浙江 杭州 310018;2.丽水学院工学院, 浙江 丽水 323000;3.浙江掌信传媒科技有限公司, 浙江 丽水 323020)

三维点云近年来在不同领域的应用越来越广泛,如人脸识别、机器人技术和生成技术等[1-5],点云特征提取方法作为以上应用开发的技术基石具有重要的研究意义。点云是表示三维对象的点集,具有稀疏性、非结构化、无序性等特性,在数据结构上与结构化的二维图像不同,这使得在计算机视觉中已经成熟的深度学习方法[6]无法直接适配到点云处理中。针对该问题,近年来出现了各种关于点云的深度学习方法。基于投影的方法[7-8]使用点云的二维投影进行推理预测,这类方法会导致信息丢失,并且其性能因投影视点变化而产生差异。基于体素的方法[9]将点云转换为体素数据结构,应用卷积进行处理,存在由于体素化而失去几何细节的问题。基于点的方法[10-11]在点云原始数据上进行研究,减少了信息的丢失。基于Transformer(一种采用自注意力机制的深度学习模型)的方法[12-15]缺乏对点云局部结构信息的利用,不能充分提取上下文语意特征且存在计算复杂度高的问题。

Transformer作为集合运算符,适合处理非结构化的、无序的点集,具有很好地适配点云数据结构的特点。因此,本文提出了一种用于点云处理的局部邻域Transformer(LNT),该方法是对点云进行邻域划分后,在局部邻域内计算自注意力,达到线性复杂度。此外,通过在自注意力计算过程中加入有效的相对位置编码,在网络构建中使用多特征融合方法,进一步提高模型性能。实验结果表明:LNT模型在ModelNet40数据集和ShapeNet数据集中表现出较好的性能优势。

2.1 设计思路

Transformer能较好地提取全局特征,但缺乏对局部特征的学习,因此研究人员针对点云数据结构进行改进,提出了一个用于三维点云处理的局部邻域Transformer。在此基础上,构建了一种新的用于点云学习的局部邻域Transformer网络。LNT的关键思想是利用Transformer固有的置换不变性,避免需要定义点云数据的顺序,引入局部邻域化的思想,通过局部自注意力机制进行特征学习。

局部邻域Transformer网络实现流程图如图1所示。首先通过最远点采样算法和KNN算法对点云进行邻域分组,然后在每个局部邻域内使用自注意力机制得到特征向量。使用金字塔式的网络结构,逐步形成全局特征表示。

图1 网络实现流程图Fig.1 Network implementation flow chart

2.2 局部邻域Transformer

给定输入为一个由N个点组成的点云P∈N×D,其中D表示特征维度。使用最远点采样和KNN对点云P进行邻域划分P={Ps∈K×D,s=1,2,3,…,S},S表示最远点采样的采样点数,K表示KNN算法中邻域内点的个数,此时P被分为S个不同的邻域(每个邻域内有K个点)。我们需要设置合理的K、S,使得不同的邻域之间能够有相互重叠的部分,进一步加强邻域间的关系。接下来,需要对邻域内的点进行自注意力的计算,由于点本身是三维坐标,减法更能表现出点与点之间的关系,并且研究人员通过实验测试了乘法和减法的效果(见本文“3.4.2”部分),最终选择使用减法计算自注意力权重:

As=σ(Q(Ps)-K(Ps)+δ)⊙(VPs)+δ)

(1)

其中,s表示第s个邻域,Q(·)、K(·)、V(·)是线性变换函数,其计算方式如公式(2)至公式(4)所示。研究人员添加了相对位置编码δ,用于补足模型对于点云结构信息的学习,σ为Softmax函数对权重进行归一化处理。

Q(Ps)=PsWq∈K×D′

(2)

K(Ps)=PsWk∈K×D′

(3)

V(Ps)=PsWv∈K×D′

(4)

其中,Wq、Wk、Wv为可学习的权重参数。

Fs=As+Ps

(5)

F=MLP(Stack(F1,F2,…,FS))∈S×D″

(6)

随后,每个邻域内的注意力特征与输入进行残差连接后,对所有邻域的输出特征(F1,F2,…,FS)进行堆叠,再经过多层感知器得到输出特征F。局部邻域Transformer结构如图2所示,其中,h表示对头注意力的头数。

图2 局部邻域TransformerFig.2 Local neighborhood Transformer

局部邻域Transformer促进了局部邻域内特征向量之间的信息交换,同时邻域间的重叠区域也促进了邻域间的信息交换。最终,所有数据点产生新的特征向量作为输出。其中,研究人员在注意力计算环节使用了多头自注意力机制[12],这让模型同时关注不同表示子空间的信息。总的来说,为了使Transformer适合点云的处理,研究人员对它进行了改进,构造了一个核心为Transformer的局部邻域自注意力层。该层集成了采样分组、局部自注意力、线性变换和一个残差连接(缓解网络退化问题)。该层能够减少特征数量,提升特征维度,为构建金字塔式的网络结构做铺垫。多层感知器中的Max Pooling运算符解决了点云置换不变性的问题。

在计算复杂度方面,由于注意力的计算集中在注意力权重分支,因此本文采用减法而非乘法,在单个邻域内的时间复杂度为O(KD),那么整体的时间复杂度则为O(SKD),优于传统注意力的二次时间复杂度,基本达到线性时间复杂度。

2.3 多特征融合

随着网络的深入,模型底层细节信息逐渐丢失,为了进一步提高模型的性能,研究人员加入了多特征融合模块。由于LNT能够减少特征个数N,提升特征维度D,因此每层输出的特征尺度都不相同,为了能够进行融合,首先需要将特征个数N进行统一,然后对特征进行合并。

第一方案是使用随机下采样的方式,使得所有底层(前两层)输出特征个数与第三层的输出特征个数统一,然后对这些特征进行合并,但模型精度并没有提升,研究人员分析原因后发现,精度与随机下采样和局部特征提取的方式有关,这导致特征之间的关系并不对应,使合并后特征变得混乱,导致最终模型效果不好。

为了解决这一问题,研究人员修改了下采样的方式,将其改为对应位置的下采样,具体来说就是由于LNT的注意力计算是将邻域内的特征聚合到中心点上的,保留了中心点的索引,因此可以通过高层的中心点索引在底层中找到对应的点,并在底层只保留这些点的特征,用于合并,多特征融合模块如图3所示。

图3 多特征融合模块Fig.3 Multi-featured fusion module

点云通过两层LNT后,点的个数逐渐减少,特征维度逐渐增加,如图3中右指向的箭头(步骤1)。为了实现特征融合,通过已知高层中心点的索引,可以找到底层特征中对应的点特征,并只保留白色的点特征(白色点),如图3左指向箭头(步骤2)。最终将每层的对应点进行合并,完成融合(步骤3)。

2.4 相对位置编码

Transformer本质上是集合算子,计算全局特征,缺乏对点云局部相对位置信息的利用。在使用自注意力处理图像的任务中,位置编码可以使得模型更好地利用数据的局部结构信息。位置编码方案很多是人工制定的,例如基于正弦和余弦函数或归一化范围值[12]。还有一些是通过学习得到的,例如在NAT[16]中,对窗口内像素位置编号后,通过可学习的偏置将位置信息加入注意力权重中,其位置偏置用于规则二维网格图像,不适用于处理离散的三维点云。

在点云处理中加入相对位置编码,可以加强模型对于点云局部结构的学习,但编码器太简单,不能充分学习结构特征,也不可设计复杂造成过拟合。

本文提出的相对位置编码,通过将两点坐标做差,引入可训练、参数化的多层感知器学习点之间的相对位置信息,使模型达到更好的效果,弥补了自注意力机制在点云处理中的缺点,解决点云的旋转不变性问题。研究人员将位置编码δ定义如下:

δ=MLP(Pi-Pj)

(7)

其中,Pi和Pj为两个点的三维坐标。编码器是线性层和激活层组成的多层感知器。相对位置编码可以强调模型更加注重提取点云三维结构上的特征,这在实验中(表4)也得到了验证:添加相对位置编码的模型与不添加的相比精度更高。因此,研究人员在公式(1)的两个分支中都添加了可训练的相对位置编码。

2.5 局部邻域Transformer网络

基于LNT,研究人员构造了完整的用于三维点云处理的网络,如图4所示。整个网络模型根据任务不同(分类任务和分割任务),网络后半段结构有所不同,但编码器的网络结构是一致的。

图4 局部邻域Transformer网络架构Fig.4 Local neighborhood Transformer network architecture

编码器:LNT能将输入点编码为一个新的高维特征向量,同时降低特征的基数。我们将它作为编码器的基础,以金字塔式的结构构建编码器。K表示KNN算法中找到邻近点的个数。将点云输入3个堆叠的LNT中,在局部邻域内学习每个点丰富的语义特征表示,然后将3层LNT的输出进行多特征融合,再经过一个全局Transformer层,最终生成输出特征。在全局Transformer层中不进行邻域的划分,因此K值设定为1。在局部邻域Transformer网络框架中,以输入2 048个点为例,编码阶段每层降采样目标点数为(512,128,64,1),每层中K值设定为(16,16,32,1)。全局自注意力是N个点与其他N-1个点计算自注意力权重,因此其空间复杂度为O(N2),而本文的局部邻域自注意力是N个点与其邻域内K个点计算自注意力权重,所以模型的参数量大幅减少,达到O(NK)。

分类:网络的细节如图4中的下分支所示。将点云数据分类为Nc个对象类别(如飞机、桌子、椅子)。将编码器输出的特征向量输入两个前馈神经网络LBRD(线性层、批归一化层和Dropout层结合的模块,其中Dropout率设定为0.5),最终由线性层预测最终分类分数,得分最高的类确定为类标签。

分割:网络的细节如图4中的上分支所示。将点云分割成Np个部件(如将飞机分割为机体、两翼、尾翼等部件),研究人员需要预测每个点的部件标签。在分割任务中采用U-Net编码器-解码器网络结构,分割网络编码器的架构与分类网络的架构相同,编码器和解码器的基本结构(网络基数和特征维度相同)是对称的,编码层的特征通过桥接到解码层辅助上采样。

图4中ULBR模块包含上采样层、线性层、归一化层和激活层。上采样层的主要功能是将特征从低分辨率的输入点集映射到高分辨率的点集上。输入特征经过三线性插值将特征映射到更高分辨率的点集上,然后将插值后的特征,与桥接的编码器阶段的特征进行结合,最终通过线性层、归一化层和激活层得到输出特征。

经由4个ULBR模块,再通过一个线性层,预测输入点云的逐点分割分数,每个点的部件标签由具有最大分数的部件标签确定,最终得到分割后的点云。

研究人员在点云分类和分割任务中评估了局部邻域Transformer网络设计的有效性,并与其他方法进行了比较。针对点云形状分类任务,在ModelNet40数据集[17]上进行测试。针对点云分割任务,在ShapeNetPart数据集[18]进行测试。

研究人员使用PyTorch实现了局部邻域Transformer网络。在实验过程中,均使用负对数似然损失函数和Adam优化器进行训练,初始学习率设置为0.001,权重衰减设置为0.000 1,训练了200个周期。

3.1 ModelNet40数据集分类实验

ModelNet40数据集包含40个类别共12 311个CAD模型,它被广泛应用于点云形状分类模型的基准测试。为保证公平性,研究人员使用官方设定的9 843个模型用于训练,2 468个模型用于测试,使用类平均精度(mAcc)和总体精度(OA)进行评估。模型在两张RTX3090显卡(24 GB显存)上运行,mini-batchsize设定为32。实验中,除了使用重采样的方式增强数据,没有使用别的方式增强数据。原则上如果附加其他训练技巧,训练精度还可以进一步提升。

实验结果如表1所示。本文提出的局部邻域Transformer在ModelNet40数据集上的总体精度(OA)为93.3%,类平均精度(mAcc)为92.0%,优于其他模型。

表 1 在ModelNet40数据集上的形状分类结果Tab.1 Shape classification results on ModelNet40 dataset

本文所提模型的类平均精度(mAcc)是最高的,达到92.0%,远超其他模型。这说明对于个别训练集数量较小的类别,本文所提模型的训练效果也很好,模型泛化性强,能够提取到丰富的特征。

3.2 ShapeNetPart数据集部件分割实验

点云部件分割是一项具有挑战性的任务,旨在将三维模型分割为多个不同的部件。本文所提模型在ShapeNetPart数据集进行了实验评估,该数据集由2 048个点的16 880个实例对象组成,共有16个类别和50个部件标签。训练中按照官方划分,14 006个对象用于训练,另外2 874个对象用于测试,其中每个实例对象分割为2—6个部件。分割任务的模型中,批处理大小、训练周期和学习率均与分类任务中的训练设置相同。

研究人员使用总体平均交并比(mIoU)作为评估指标,并列举出了每个类平均交并比。实验中,将本文所提模型与其他模型(如PointNet、PointNet++、M-Conv等)进行了比较,每个类别中的最高精度用加粗字体标识出来,结果如表2所示。结果表明,LNT模型的总体分割精度实现了超越,类别精度在包、汽车、刀等7个类别上取得了最佳结果。图5展示了分割效果可视化实例,第一行为未进行分割处理的对象,第二行为LNT模型分割预测的结果。

表 2 在ShapeNetPart数据集上的部件分割结果Tab.2 Part segmentation results on ShapeNetPart dataset

图5 ShapeNet数据集分割可视化实例Fig.5 Example of ShapeNet dataset segmentation visualization

3.3 模型空间复杂度

当前,使用自注意力机制的其他模型大多使用的全局自注意力,导致这些模型在注意力部分的空间复杂度为O(N2)。相比之下,本文所提的LNT模型在局部邻域内计算自注意力,可以让模型的参数量大幅减少,达到O(NK2),其中N为点的个数,K为邻域内点的个数(K<32)。

表3以分割任务为例,展示了LNT模型的参数量,并与其他模型进行了对比,表明LNT在空间复杂度上的优势。

表 3 模型参数量对比Tab.3 Comparison of model parameters

3.4 消融性实验

3.4.1 相对位置编码的影响

通过引入可训练的、参数化的位置编码器来学习点与点之间的相对位置信息,可以使得局部自注意力机制更好地利用点云的结构信息;并以ModelNet40数据集上的分类任务为例,具体研究了相对位置编码添加位置对模型的影响,结果如表4所示。

表 4 添加相对位置编码的实验对比Tab.4 Experimental comparison of adding relative position coding

从表4中可以看到:如果没有位置编码,模型的精度就会显著下降。采用相对位置编码时的模型精度高于无位置编码。可见,相对位置编码可以提高网络模型的性能。当相对位置编码只添加到注意力权重计算分支或仅添加到特征转换分支时,模型精度会不如整体都添加时的模型精度。可见,在两个分支都添加相对位置编码是很重要的。

3.4.2 自注意力计算方式的影响

Transformer中自注意力权重计算的计算方式有多种,如减法、乘法等,其中乘法(WqTWk)是Transformer通用的计算方式。但是,研究人员认为减法的计算方式(Wq-Wk)要更适合点云自注意力的计算,因为点云本身是三维空间中的位置信息,减法能更好地提取点与点之间的位置关系,同时更能突出点与点之间的差异性,更有利于特征的学习。研究人员通过实验证明了这一点,使用乘法计算注意力权重的总体精度为92.5%,与使用减法计算注意力权重的总体精度93.3%相比,低了0.8%。

3.4.3K值的设定

超参数K表示最近邻采样点的数量即邻域的范围大小,K值的合理设定是决定模型性能的重要一环,研究人员需要设定合适的K值保证邻域之间能够有相互重叠的部分,以此进一步加强邻域间的关系。在网络的底层,研究人员希望模型专注于低级特征的学习,因此在前两层设定的K值较小;在网络的高层,设定较大的K值对特征进行聚合,形成全局特征。

可以通过实验寻找每层合适的K值设定,结果见表5所示。最合理的每层K值设定为(16,16,32,1)。

表 5 K值的选择Tab.5 Selection of K value

本文提出了一种用于点云特征提取的局部邻域Transformer模型,该模型从点云数据结构的特点出发,利用Transformer自注意力机制在点云局部几何结构中学习特征,能够充分提取上下文语意信息,提升特征提取能力,并达到线性计算复杂度。同时,多特征融合方法能有效地利用模型底层特征,进一步提升模型的性能。实验结果表明,LNT模型在ModelNet40数据集形状分类任务中的总体精度可达到93.3%,类平均精度可达到92.0%,优于现有其他模型,并且在分割任务中点云分割结果也是有效的。同时,实验证明了模型中的相对位置编码、局部自注意力计算方式及K值设定的合理性。未来,研究人员会对模型的网络架构进行改进,在点云分割、语义分割、点云生成等任务中做进一步的扩展应用。