基于局部鄰域Transformer的點云特征提取方法

張海博, 沈 洋,, 許 浩, 包艷霞,3, 劉 江

(1.浙江理工大學計算機科學與技術學院, 浙江 杭州 310018;2.麗水學院工學院, 浙江 麗水 323000;3.浙江掌信傳媒科技有限公司, 浙江 麗水 323020)

1 引言(Introduction)

三維點云近年來在不同領域的應用越來越廣泛,如人臉識別、機器人技術和生成技術等[1-5],點云特征提取方法作為以上應用開發的技術基石具有重要的研究意義。點云是表示三維對象的點集,具有稀疏性、非結構化、無序性等特性,在數據結構上與結構化的二維圖像不同,這使得在計算機視覺中已經成熟的深度學習方法[6]無法直接適配到點云處理中。針對該問題,近年來出現了各種關于點云的深度學習方法。基于投影的方法[7-8]使用點云的二維投影進行推理預測,這類方法會導致信息丟失,并且其性能因投影視點變化而產生差異。基于體素的方法[9]將點云轉換為體素數據結構,應用卷積進行處理,存在由于體素化而失去幾何細節的問題。基于點的方法[10-11]在點云原始數據上進行研究,減少了信息的丟失。基于Transformer(一種采用自注意力機制的深度學習模型)的方法[12-15]缺乏對點云局部結構信息的利用,不能充分提取上下文語意特征且存在計算復雜度高的問題。

Transformer作為集合運算符,適合處理非結構化的、無序的點集,具有很好地適配點云數據結構的特點。因此,本文提出了一種用于點云處理的局部鄰域Transformer(LNT),該方法是對點云進行鄰域劃分后,在局部鄰域內計算自注意力,達到線性復雜度。此外,通過在自注意力計算過程中加入有效的相對位置編碼,在網絡構建中使用多特征融合方法,進一步提高模型性能。實驗結果表明:LNT模型在ModelNet40數據集和ShapeNet數據集中表現出較好的性能優勢。

2 方法設計(Method design)

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模塊,再通過一個線性層,預測輸入點云的逐點分割分數,每個點的部件標簽由具有最大分數的部件標簽確定,最終得到分割后的點云。

3 實驗和結果(Experiment and results)

研究人員在點云分類和分割任務中評估了局部鄰域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

4 結論(Conclusion)

本文提出了一種用于點云特征提取的局部鄰域Transformer模型,該模型從點云數據結構的特點出發,利用Transformer自注意力機制在點云局部幾何結構中學習特征,能夠充分提取上下文語意信息,提升特征提取能力,并達到線性計算復雜度。同時,多特征融合方法能有效地利用模型底層特征,進一步提升模型的性能。實驗結果表明,LNT模型在ModelNet40數據集形狀分類任務中的總體精度可達到93.3%,類平均精度可達到92.0%,優于現有其他模型,并且在分割任務中點云分割結果也是有效的。同時,實驗證明了模型中的相對位置編碼、局部自注意力計算方式及K值設定的合理性。未來,研究人員會對模型的網絡架構進行改進,在點云分割、語義分割、點云生成等任務中做進一步的擴展應用。