多尺度融合卷積的輕量化Transformer無(wú)人機(jī)地物識(shí)別模型

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

(西南石油大學(xué) 計(jì)算機(jī)科學(xué)學(xué)院 四川 成都 610500)

0 引言

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

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

1 相關(guān)工作

1.1 無(wú)人機(jī)遙感圖像語(yǔ)義分割

遙感圖像被廣泛應(yīng)用于眾多領(lǐng)域,如車道線檢測(cè)[11]、自動(dòng)化農(nóng)業(yè)[12]等。傳統(tǒng)的遙感圖像語(yǔ)義分割方式大多依賴圖像的光譜信息和紋理信息。然而,不同的地物可能具有相同的光譜特征或者紋理特征,相同的地物也可能具有不同的光譜特征或者紋理特征[7],這導(dǎo)致傳統(tǒng)分割方法的精度不足。而基于深度學(xué)習(xí)的語(yǔ)義分割方式不僅會(huì)考慮地物的光譜信息和紋理信息,還會(huì)考慮地物間相互關(guān)系,這使得遙感圖像地物識(shí)別在效率和精度上得到提升。

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

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

1.2 用于圖像的Transformer網(wǎng)絡(luò)

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

1) 在拉直為序列的過程中,圖片塊內(nèi)部的空間結(jié)構(gòu)遭到了破壞。

2) 當(dāng)ViT用于大圖像時(shí),由于劃分的圖片塊數(shù)量增加以及自注意力計(jì)算具有二次復(fù)雜度,導(dǎo)致計(jì)算成本相當(dāng)高。

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

1.3 Transformer在圖像領(lǐng)域中的輕量化方法

SETR自注意力計(jì)算的復(fù)雜度為O(n2),表明模型的計(jì)算復(fù)雜度會(huì)隨著圖像尺寸的增大而快速增加,這對(duì)于資源受限的無(wú)人機(jī)遙感任務(wù)并不友好。降低自注意力計(jì)算的復(fù)雜度是當(dāng)前的一個(gè)研究熱點(diǎn)。文獻(xiàn)[19]中PVT引入SRA模塊,通過卷積的方式縮減鍵值序列,減少了自注意力計(jì)算的復(fù)雜度。PVTv2[20]在PVT的基礎(chǔ)上,將原有的SRA模塊的卷積操作換成平均池化操作,進(jìn)一步減少了自注意力計(jì)算的復(fù)雜度。Swin Transformer[21]通過新穎的分窗以及移窗操作,將原本基于全局的自注意力計(jì)算變?yōu)榫植孔宰⒁饬τ?jì)算,可以將二次方的復(fù)雜度降為線性復(fù)雜度。

2 問題描述和整體框架

2.1 問題描述

本文通過語(yǔ)義分割的方式進(jìn)行地物識(shí)別,向網(wǎng)絡(luò)中輸入原圖,利用訓(xùn)練后的模型進(jìn)行推理,從而得到與標(biāo)簽接近的預(yù)測(cè)圖。模型的輸入為X∈RH×W×3,輸出為Y∈RH×W×Ncls。其中:H代表圖片的高;W代表圖片的寬;Ncls代表類別數(shù)量。

2.2 整體框架

模型整體結(jié)構(gòu)如圖1所示。采用雙分支端到端的網(wǎng)絡(luò)結(jié)構(gòu),上、下分支同時(shí)進(jìn)行計(jì)算。

圖1 模型整體結(jié)構(gòu)Figure 1 Overall structure of the model

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

3.1 多尺度融合卷積

由于只使用Transformer模塊會(huì)丟失圖片塊原有的內(nèi)部空間信息,使用預(yù)訓(xùn)練后的ResNet34來(lái)學(xué)習(xí)圖片塊丟失的內(nèi)部空間結(jié)構(gòu)信息。Transformer是直接基于全局語(yǔ)義相關(guān)性構(gòu)建的模型,因此該模型能夠獲取足夠大的有效感受野。為了減少模型參數(shù)以及盡可能保留Transformer獲取的感受野,簡(jiǎn)單堆疊的MLP層被用來(lái)對(duì)特征圖進(jìn)行細(xì)微調(diào)整。同時(shí),為了獲取多尺度信息,模型融合了4個(gè)階段的特征圖。

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

對(duì)于單個(gè)注意力模塊,如圖2所示,采用類似Segformer[23]的Transformer架構(gòu)。

圖2 Transformer模塊Figure 2 Transformer module

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

假設(shè)頭部索引為i,有

Vi=Vi+LE(Vi),

(1)

(2)

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

XSSA=SSA(OverlapPatchEmbed(Xk)),

Xout1=MLP(XSSA),

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

Xk+1=MLP(Xout2)+XSSA,

(3)

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

3.3 多尺度MLP上采樣模塊

圖3 MLP解碼器Figure 3 MLP decoder

4 實(shí)驗(yàn)設(shè)置

4.1 數(shù)據(jù)集

實(shí)驗(yàn)使用國(guó)際攝影測(cè)量與遙感協(xié)會(huì)(ISPRS)的Vaihingen和Potsdam數(shù)據(jù)集,其數(shù)據(jù)都是由無(wú)人機(jī)在城市上空進(jìn)行多次拍攝獲得的,拍攝地點(diǎn)為德國(guó)的Vaihingen和Potsdam兩個(gè)城市。數(shù)據(jù)集包含6種常見的土地覆蓋類別,分別為樹木、汽車、建筑物、低矮植被、不透水平面和其他。

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

Potsdam數(shù)據(jù)集包含38張圖片,其中24張圖片具備標(biāo)注信息,每張圖片的像素大小為6 000×6 000,地面采樣距離為5 cm。將18張圖片作為訓(xùn)練集,6張圖片作為測(cè)試集。將所有圖片切分為256×256像素大小,使用隨機(jī)水平和垂直翻轉(zhuǎn)數(shù)據(jù)增強(qiáng)操作。

4.2 評(píng)價(jià)指標(biāo)

采用ISPRS的官方評(píng)估方法,所用評(píng)價(jià)指標(biāo)為整體準(zhǔn)確度(OA)、F1值以及交并比(IoU)。

OA的計(jì)算公式為

(4)

其中:TP、TN、FP、FN分別代表真陽(yáng)性、真陰性、假陽(yáng)性、假陰性。

F1值的計(jì)算公式為

(5)

(6)

(7)

IoU的計(jì)算公式為

(8)

其中:Np表示預(yù)測(cè)集;Ngt表示圖像的標(biāo)簽。

4.3 訓(xùn)練設(shè)置

實(shí)驗(yàn)通過pytorch1.10.1、python3.8實(shí)現(xiàn),所有實(shí)驗(yàn)通過一張NVIDIA RTX 3090顯卡進(jìn)行訓(xùn)練。網(wǎng)絡(luò)使用AdamW優(yōu)化器,學(xué)習(xí)率為6×10-5,權(quán)重衰減為0.01。同時(shí),采用power為0.9的“poly”學(xué)習(xí)率策略,迭代次數(shù)為3.2×105。對(duì)比算法均采用相同的訓(xùn)練策略。

5 實(shí)驗(yàn)結(jié)果

將本文模型與近幾年具有代表性的模型以及基于ViT進(jìn)行分割的模型進(jìn)行比較,不同模型在Vaihingen和Potsdam數(shù)據(jù)集上的對(duì)比實(shí)驗(yàn)結(jié)果分別如表1、表2所示。其中,對(duì)最優(yōu)結(jié)果進(jìn)行了加黑顯示。

表2 不同模型在Potsdam數(shù)據(jù)集上的對(duì)比實(shí)驗(yàn)結(jié)果Table 2 Comparative experimental results of different models on Potsdam data set 單位:%

1) 定量比較

從表1、表2可以看出,在Vaihingen數(shù)據(jù)集上,本文模型的F1、IoU平均值分別比對(duì)比模型中的最優(yōu)結(jié)果高出0.12%、0.02%。在Potsdam數(shù)據(jù)集上,本文模型的F1、IoU平均值以及OA值分別比對(duì)比模型中的最優(yōu)結(jié)果高出0.27%、0.39%、0.29%。

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

主要有以下幾個(gè)方面導(dǎo)致SETR的表現(xiàn)不佳:① SETR特征圖大小沒有發(fā)生變化,對(duì)于粗粒度和細(xì)粒度都以同一尺度特征圖進(jìn)行輸出,難以獲取粗粒度和細(xì)粒度的特征。② SETR未使用跨步卷積,而是直接切分成小塊,切斷了SETR中塊之間的聯(lián)系。③ SETR將切分好的圖像塊進(jìn)行拉直操作,破壞了圖像塊的內(nèi)部空間結(jié)構(gòu),對(duì)于小目標(biāo)的識(shí)別十分不友好。

2) 定性比較

本文方法在類別較為平衡的數(shù)據(jù)集Postdam上取得了較為理想的結(jié)果,同時(shí)在不平衡的數(shù)據(jù)集Vaihingen上也取得了總體不錯(cuò)的效果。從表1可以看出,本文方法在一定程度上緩解了類別不均衡的問題。對(duì)于數(shù)量較多的建筑物類別,本文方法相較基于CNN的方法也有明顯的優(yōu)勢(shì),這是由于通過雙分支的方法能夠?qū)ransformer和CNN的優(yōu)點(diǎn)進(jìn)行有機(jī)融合。

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

圖4 Vaihingen數(shù)據(jù)集實(shí)驗(yàn)結(jié)果可視化Figure 4 Visualization of experimental results on Vaihingen data set

圖5 Potsdam數(shù)據(jù)集實(shí)驗(yàn)結(jié)果可視化Figure 5 Visualization of experimental results on Potsdam data set

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

6 結(jié)語(yǔ)

本文提出一種多尺度融合卷積的輕量化Transformer模型,實(shí)現(xiàn)了Transformer與卷積網(wǎng)絡(luò)的有機(jī)融合。利用改進(jìn)后的Transformer獲取了足夠大的

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

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