多尺度融合卷積的輕量化Transformer無人機地物識別模型

肖 斌, 羅 浩, 張恒賓, 劉宏偉, 張興鵬

(西南石油大學 計算機科學學院 四川 成都 610500)

0 引言

與高空衛星遙感影像相比,無人機航拍影像具有成像分辨率高、大氣因素干擾少等優點[1],更易區分建筑、汽車、樹木等細小物體。因此,無人機航拍影像在小區域遙感應用方面具有很大的前景。在遙感圖像中,語義分割方法能夠很好地確定土地類型,提取建筑和道路信息。隨著深度學習的發展,基于卷積神經網絡(convolutional neural network, CNN)的全卷積神經網絡(fully convolution network, FCN)[2]實現了像素級分類。SegNet[3]、UNet[4]等一系列網絡進一步提高了語義分割的精度。Zhao等[5]使用無人機在100 m高度采集圖像,并使用UNet識別小麥倒伏。然而,同一類地物可能具有不同的形狀和大小,不同類地物也可能具有相似的特征[6],只利用光譜信息的模型不足以有效區分地物,還需要通過多個尺度的上下文信息進行輔助識別。由于FCN的卷積核感受野大小固定,不能獲取足夠的上下文信息,導致語義分割網絡的分割精度不足。吳蘭等[7]通過多尺度融合的方式有效地獲取了上下文信息。

Vision Transformer(ViT)[8]舍棄了卷積運算,采用純Transformer[9]架構對全局語義信息進行建模,使得模型有更大的感受野。然而,ViT在處理高分辨率圖像時往往會產生巨大的參數量。例如,利用SETR[10]處理像素大小256×256的圖片,參數量達到驚人的317 M?；诖?本文提出一種用于無人機遙感圖像的多尺度融合卷積的輕量化Transformer模型,旨在研究Transformer在無人機遙感圖像語義分割上的潛力。主要貢獻如下:① 設計了一種輕量級的多尺度融合卷積方法,使得Transformer在計算自注意力時丟失的塊內空間信息能夠通過CNN得到補償。② 設計了多尺度縮減鍵值序列的方式,可以有效減少自注意力計算的參數量,同時捕獲遠程、小型孤立對象。③ 設計了輕量級的多尺度MLP解碼器,丟棄了不必要的卷積操作。實驗結果表明,所提模型在Vaihingen和Potsdam數據集上性能有所提升。

1 相關工作

1.1 無人機遙感圖像語義分割

遙感圖像被廣泛應用于眾多領域,如車道線檢測[11]、自動化農業[12]等。傳統的遙感圖像語義分割方式大多依賴圖像的光譜信息和紋理信息。然而,不同的地物可能具有相同的光譜特征或者紋理特征,相同的地物也可能具有不同的光譜特征或者紋理特征[7],這導致傳統分割方法的精度不足。而基于深度學習的語義分割方式不僅會考慮地物的光譜信息和紋理信息,還會考慮地物間相互關系,這使得遙感圖像地物識別在效率和精度上得到提升。

目前,已有許多深度學習語義分割方法應用在遙感圖像上。例如,Chong等[13]提出CFEM和IEM模塊,增強了卷積網絡區分小尺度物體和細化小尺度物體邊界的能力。Ding等[14]使用輕量級的PAM和AEM模塊,平衡了模型的特征表示能力和空間定位精度。然而,通過這些模塊收集的上下文信息仍然不夠,導致一些較小目標的分割結果并不太理想。

現有的語義分割方法大多集中在擴大有效感受野或者通過上下文建模的方式來獲取足夠的上下文信息。文獻[15]中Deeplabv3+引入空洞卷積,通過擴大卷積核來擴大感受野。文獻[16]中PSPNet引入PPM模塊,通過上下文建模來獲取足夠的上下文信息。近年來,利用注意力機制來捕獲遠程上下文信息的方式開始興起。文獻[17]中DANet引入通道和空間注意力,提升了模型性能。

1.2 用于圖像的Transformer網絡

相比CNN,ViT這種基于全局語義相關性構建的模型在遠距離特征捕獲能力上具有更大的優勢[18]。SETR將ViT作為語義分割模型的編碼器,ViT將圖片劃分成像素大小為16×16的圖片塊,再將其拉直為序列,送入Transformer進行注意力計算。然而,ViT存在以下問題:

1) 在拉直為序列的過程中,圖片塊內部的空間結構遭到了破壞。

2) 當ViT用于大圖像時,由于劃分的圖片塊數量增加以及自注意力計算具有二次復雜度,導致計算成本相當高。

3) ViT作為序列到序列的模型,無法獲得與常規CNN類似的多尺度特征圖,這意味著其不能很好地獲取圖片粗粒度與細粒度的特征信息。

1.3 Transformer在圖像領域中的輕量化方法

SETR自注意力計算的復雜度為O(n2),表明模型的計算復雜度會隨著圖像尺寸的增大而快速增加,這對于資源受限的無人機遙感任務并不友好。降低自注意力計算的復雜度是當前的一個研究熱點。文獻[19]中PVT引入SRA模塊,通過卷積的方式縮減鍵值序列,減少了自注意力計算的復雜度。PVTv2[20]在PVT的基礎上,將原有的SRA模塊的卷積操作換成平均池化操作,進一步減少了自注意力計算的復雜度。Swin Transformer[21]通過新穎的分窗以及移窗操作,將原本基于全局的自注意力計算變為局部自注意力計算,可以將二次方的復雜度降為線性復雜度。

2 問題描述和整體框架

2.1 問題描述

本文通過語義分割的方式進行地物識別,向網絡中輸入原圖,利用訓練后的模型進行推理,從而得到與標簽接近的預測圖。模型的輸入為X∈RH×W×3,輸出為Y∈RH×W×Ncls。其中:H代表圖片的高;W代表圖片的寬;Ncls代表類別數量。

2.2 整體框架

模型整體結構如圖1所示。采用雙分支端到端的網絡結構,上、下分支同時進行計算。

圖1 模型整體結構Figure 1 Overall structure of the model

3 多尺度融合卷積的輕量化Transformer模型

3.1 多尺度融合卷積

由于只使用Transformer模塊會丟失圖片塊原有的內部空間信息,使用預訓練后的ResNet34來學習圖片塊丟失的內部空間結構信息。Transformer是直接基于全局語義相關性構建的模型,因此該模型能夠獲取足夠大的有效感受野。為了減少模型參數以及盡可能保留Transformer獲取的感受野,簡單堆疊的MLP層被用來對特征圖進行細微調整。同時,為了獲取多尺度信息,模型融合了4個階段的特征圖。

3.2 縮減鍵值序列的自注意力Transformer

對于單個注意力模塊,如圖2所示,采用類似Segformer[23]的Transformer架構。

圖2 Transformer模塊Figure 2 Transformer module

由于原有SRA只在單一尺度上對鍵值序列進行縮減,這并不利于小物體的分割,使用分流的自注意力(shunted self-attention, SSA)[24]進行改進。SSA操作使用多尺度縮減鍵值序列的方式獲取目標多個尺度的上下文相互關系,通過頭部進行分組,不同的頭部鍵值序列的縮減尺度不同。

假設頭部索引為i,有

Vi=Vi+LE(Vi),

(1)

(2)

其中:dh表示維度。為了最大限度地輕量化模型,保留了Segformer原有的Mix-FFN設計。Transformer模塊的整體計算過程為

XSSA=SSA(OverlapPatchEmbed(Xk)),

Xout1=MLP(XSSA),

Xout2=GELU(Conv3×3(Xout1)),

Xk+1=MLP(Xout2)+XSSA,

(3)

其中:Xk表示第k個輸入特征圖;XSSA表示經過SSA計算后的輸出;Xk+1表示Transformer模塊的輸出;OverlapPatchEmbed(·)表示通過跨步卷積對圖像進行分塊;MLP(·)表示全連接層;Conv3×3(·)表示3×3的深度卷積;GELU(·)表示激活函數[25]。

3.3 多尺度MLP上采樣模塊

圖3 MLP解碼器Figure 3 MLP decoder

4 實驗設置

4.1 數據集

實驗使用國際攝影測量與遙感協會(ISPRS)的Vaihingen和Potsdam數據集,其數據都是由無人機在城市上空進行多次拍攝獲得的,拍攝地點為德國的Vaihingen和Potsdam兩個城市。數據集包含6種常見的土地覆蓋類別,分別為樹木、汽車、建筑物、低矮植被、不透水平面和其他。

Vaihingen數據集包含33張圖片,其中16張圖片具備標注信息,圖片的平均像素大小為2 494×2 064,地面采樣距離為9 cm。將12張圖片作為訓練集,4張圖片作為測試集。此外,每幅圖像包含近紅外波段、紅色波段、綠色波段三種波段。

Potsdam數據集包含38張圖片,其中24張圖片具備標注信息,每張圖片的像素大小為6 000×6 000,地面采樣距離為5 cm。將18張圖片作為訓練集,6張圖片作為測試集。將所有圖片切分為256×256像素大小,使用隨機水平和垂直翻轉數據增強操作。

4.2 評價指標

采用ISPRS的官方評估方法,所用評價指標為整體準確度(OA)、F1值以及交并比(IoU)。

OA的計算公式為

(4)

其中:TP、TN、FP、FN分別代表真陽性、真陰性、假陽性、假陰性。

F1值的計算公式為

(5)

(6)

(7)

IoU的計算公式為

(8)

其中:Np表示預測集;Ngt表示圖像的標簽。

4.3 訓練設置

實驗通過pytorch1.10.1、python3.8實現,所有實驗通過一張NVIDIA RTX 3090顯卡進行訓練。網絡使用AdamW優化器,學習率為6×10-5,權重衰減為0.01。同時,采用power為0.9的“poly”學習率策略,迭代次數為3.2×105。對比算法均采用相同的訓練策略。

5 實驗結果

將本文模型與近幾年具有代表性的模型以及基于ViT進行分割的模型進行比較,不同模型在Vaihingen和Potsdam數據集上的對比實驗結果分別如表1、表2所示。其中,對最優結果進行了加黑顯示。

表2 不同模型在Potsdam數據集上的對比實驗結果Table 2 Comparative experimental results of different models on Potsdam data set 單位:%

1) 定量比較

從表1、表2可以看出,在Vaihingen數據集上,本文模型的F1、IoU平均值分別比對比模型中的最優結果高出0.12%、0.02%。在Potsdam數據集上,本文模型的F1、IoU平均值以及OA值分別比對比模型中的最優結果高出0.27%、0.39%、0.29%。

如表1所示,除了SETR外,基于ViT的方法在OA值上普遍高于基于CNN的方法。這是由于通過Transformer的方式獲取了足夠的感受野,使得原來難以確定的邊界通過其上下文信息能夠進行較為精準的判斷。同時,基于ViT的方法在類別數量較多時預測會更加準確?？梢钥吹?在不平衡的數據集Vaihingen上,對于類別較多的建筑物,Swin Transformer達到了最好的效果;對于類別較少的汽車,反而是PSPNet達到了最好的效果。這是由于ViT將塊拉直為序列的操作導致了這樣的現象。建筑物類別具有數量多且尺度大等特點,拉直成序列后對于建筑物的影響并不是特別大。而汽車類別在劃分的圖像塊中占比較小,尺度也較小,拉直成序列后塊內結構信息損失嚴重,因此會更加難以判斷。

主要有以下幾個方面導致SETR的表現不佳:① SETR特征圖大小沒有發生變化,對于粗粒度和細粒度都以同一尺度特征圖進行輸出,難以獲取粗粒度和細粒度的特征。② SETR未使用跨步卷積,而是直接切分成小塊,切斷了SETR中塊之間的聯系。③ SETR將切分好的圖像塊進行拉直操作,破壞了圖像塊的內部空間結構,對于小目標的識別十分不友好。

2) 定性比較

本文方法在類別較為平衡的數據集Postdam上取得了較為理想的結果,同時在不平衡的數據集Vaihingen上也取得了總體不錯的效果。從表1可以看出,本文方法在一定程度上緩解了類別不均衡的問題。對于數量較多的建筑物類別,本文方法相較基于CNN的方法也有明顯的優勢,這是由于通過雙分支的方法能夠將Transformer和CNN的優點進行有機融合。

對Vaihingen和Potsdam數據集的部分實驗結果進行了可視化,分別如圖4、圖5所示?？梢钥闯?對于Vaihingen數據集,本文方法對陰影部分具有很好的處理能力,這得益于Transformer與CNN的有機結合獲取了足夠大的感受野;對于Potsdam數據集,在原圖上光譜信息相似的地方(例如第2張圖中的樹木和低矮植被),由于能夠有效獲取上下文信息,模型能夠較好地對這些地方進行分割。

圖4 Vaihingen數據集實驗結果可視化Figure 4 Visualization of experimental results on Vaihingen data set

圖5 Potsdam數據集實驗結果可視化Figure 5 Visualization of experimental results on Potsdam data set

不同模型在Vaihingen和Potsdam數據集上的參數量如圖6所示?？梢钥闯?本文模型是使用Transformer的網絡模型中參數量最少且效果最好的。在Potsdam數據集上,與雙分支網絡STransFuse的參數量(86.99 M)相比,本文模型的參數量為71.24 M,減少18%,這得益于SSA以及MLP解碼器。

6 結語

本文提出一種多尺度融合卷積的輕量化Transformer模型,實現了Transformer與卷積網絡的有機融合。利用改進后的Transformer獲取了足夠大的

圖6 不同模型在Vaihingen和Potsdam數據集上的參數量Figure 6 Parameters of different models on Vaihingen and Potsdam data sets

感受野,同時又能結合卷積網絡的優點,這使得模型在語義分割任務上獲得了較大的優勢。為了進一步補充圖像粗粒度和細粒度特征,設計了多尺度縮減鍵值序列的注意力計算方式,加強了網絡對于小目標的識別能力。由于無人機遙感任務通常硬件資源受限,采用了輕量級的多尺度MLP解碼器網絡以及添加了輕量級的MLP層來幫助融合Transformer和卷積網絡。通過對比實驗可知,本文模型在精度上相比其他模型有較大的優勢。并且,本文模型相比基于Transformer的模型有更少的參數量。Transformer在遙感任務上擁有很大的潛力,未來將專注于輕量化Transformer,使其能夠更加適應無人機遙感任務。