大規模視覺基礎模型在基于自然圖像的視覺任務中取得了重大進展。得益于良好的可擴展性和表征能力,基于視覺Transformer (Vision Transformer, ViT) 的大規模視覺基礎模型吸引了研究社區的廣泛關注,并在多種視覺感知任務中廣泛應用。然而,在遙感圖像感知領域,大規模視覺模型的潛力尚未得到充分的探索。為此,我們首次提出面向遙感任務設計的大規模視覺基礎模型[1],針對具有1億參數的一般結構的ViT (Plain ViT),設計符合遙感圖像特點的新型注意力機制,并據此研究大模型在遙感圖像理解任務上的性能,包括圖像分類、目標檢測、語義分割等。
簡單來說,為了更好地應對遙感圖像所具有的大尺寸特性以及圖像中的目標物體的朝向任意性,我們提出了一種新的旋轉可變尺寸窗口的注意力(Rotated Varied-Size Window Attention, RVSA)來代替Transformer中的原始完全注意力(Vanilla Full Self-Attention),它可以從生成的不同窗口中提取豐富的上下文信息來學習更好的目標表征,并顯著降低計算成本和內存占用。
實驗表明,在檢測任務上,我們提出的模型優于目前為止所有最先進的模型,其在DOTA-V1.0數據集上取得了81.24% mAP的最高精度。在下游分類和分割任務上,所提出的模型與現有先進方法相比性能具有很好的競爭力。進一步的分析實驗表明該模型在計算復雜度、遷移學習的樣本效率、可解釋性等方面具有明顯優勢。
本工作由京東探索研究院、武漢大學以及悉尼大學聯合完成,已被IEEE TGRS接收。
01
研究背景
在遙感圖像感知領域中,卷積神經網絡(Convolutional Neural Network, CNN)是提取多尺度視覺特征最常用的模型。然而,卷積操作的感受野受限,這使得CNN很難關注長距離像素并提取全局上下文信息。為了解決這一問題,研究者提出使用自注意力(Self-Attention, SA)機制,通過計算圖像中任意像素(特征)之間的相似性來靈活地建模特征之間的長距依賴關系。這一技術在計算機視覺領域的諸多任務上取得了良好的表現。其中,視覺Transformer模型采用了多頭自注意力(Multi-Head Self-Attention, MHSA)的設計,在多個投影子空間中同時計算自注意力,使得提取的上下文信息更加多樣化,從而進一步提高了特征的表征能力。
最早提出的視覺Transformer模型ViT [2]的結構采用了非層次化的一般結構設計,即在特征嵌入層之后重復堆疊Transformer編碼器模塊,其中每個模塊輸出的空間尺度、特征維度均相同。為了更好地使ViT適應下游任務,研究人員借用了CNN中的分層設計思想,并相應地設計了層次化視覺Transformer[3, 4]。這些模型通常使用大規模數據集并以有監督的方式進行預訓練,然后再在下游任務的訓練集上進行微調。最近,探索研究院通過比較不同的預訓練方法和模型,將層次化視覺Transformer應用于遙感圖像上并對其性能進行了詳細的實證研究[5],驗證了層次化視覺Transformer相比于CNN的優勢以及使用大規模遙感場景標注數據集進行預訓練的有效性。然而,是否一定要采用層次化結構的模型才能在遙感圖像上獲得較好性能呢?在本項研究中,我們首次嘗試采用非層次化結構的模型并驗證了其在一系列遙感圖像感知任務上的優勢和潛力。
具體來說,我們首先使用具有約一億參數的Plain ViT模型和研究院最近提出的更先進的ViTAE 模型[6],并采用掩碼圖像建模算法MAE [7]在大規模遙感數據集MillionAID [8]上對其進行預訓練,從而得到很好的初始化參數。
在預訓練完成后,我們通過在下游任務相關數據集上進行微調,從而完成相應任務。由于下游任務的圖像分辨率較大,為了降低視覺Transformer在下游任務上的計算成本和內存占用,研究者通常采用窗口注意力(Window-based Attention)機制來代替原始的完全注意力機制。然而,窗口注意力采用的固定窗口大小和位置會限制模型提取上下文信息的范圍以及跨窗信息交互,從而影響模型的表征能力。
為此,探索研究院提出了一種名為可變大小窗口的注意力機制(Varied-Size Window Attention, VSA) [9]。它通過學習窗口的縮放和偏移因子,以使窗口的大小、形狀和位置適應不同的圖像內容,從而提高特征的表征能力,在多個視覺感知任務中獲得了更好的性能。不同于自然圖像中目標主要呈現上下方向的特點,遙感圖像中的目標具有任意朝向,如圖1所示。為了處理這種差異,我們進一步引入了一種可學習的旋轉框機制,從而獲得具有不同角度、大小、形狀和位置的窗口,實現了提取更豐富的上下文新型的目標。
圖1:兩種常見類別(橋梁和飛機)的自然圖像(a)與遙感圖像(b)的區別
基于ViT和ViTAE模型,我們將上述自注意力方法應用于三種遙感感知任務(場景分類、語義分割和目標檢測),并開展了詳細的實驗評估,取得了很好的效果。我們希望這項研究能夠填補遙感大模型領域的空白,并為遙感社區發展更大規模的Plain ViT模型提供有益的參考。
02
方法介紹
2.1 MillionAID
MillionAID [8]是一個具有遙感場景圖像和標簽的大型數據集。它包含1,000,848個RGB格式的非重疊遙感場景,非常適合用于深度神經網絡模型預訓練。該數據集包含51類,每類有大約2,000-45,000個圖像。該數據集中的圖片是從包含各種傳感器和不同分辨率數據的谷歌地球上收集得到的。圖像尺寸分布廣泛,覆蓋了110*110到31,672*31,672個像素的多種情況。應該注意的是,盡管該數據集同時包含圖像和標簽,但在本項研究中,我們只采用圖像數據進行無監督預訓練。
2.2 MAE
MAE [7]是一種生成式自監督預訓練方法,采用了非對稱的網絡結構提取非掩碼區域的圖像特征并預測掩碼區域的圖像內容,具有很高的計算效率。它首先將圖像分割成不重疊的圖像塊,然后通過特征嵌入層將每個圖像塊映射為視覺Token。按照一定掩碼比率,一些Token被刪除并被作為要預測的掩碼區域。剩余的Token被饋送到Transformer編碼器網絡進行特征提取。然后,解碼器部分利用編碼器提取到的可見區域Token的特征和可學習的掩碼區域的Token來恢復掩碼區域圖像內容。在訓練過程中,通過最小化像素空間或特征空間中掩碼區域的預測和圖像真值之間的差異來訓練模型。我們遵循原始MAE文獻中的設置并在歸一化像素空間中計算訓練損失。
2.3 MAE無監督預訓練
圖2:預訓練階段的視覺Transformer的網絡結構
圖2展示了所采用的Plain ViT模型的基本結構。具體來說,我們采用兩種骨干網絡ViT和ViTAE進行預訓練。前者由具有完全自注意力的Plain ViT編碼器組成。這種簡單的結構能夠使其無縫地采用MAE方法進行預訓練。相比之下,ViTAE引入了卷積結構從而讓網絡獲取局部性歸納偏置,即采用與MHSA并列的平行卷積分支PCM。在預訓練時,因為MAE中的隨機掩蔽策略破壞了空間關系,我們將PCM的卷積和從3*3改為1*1,以避免其學習到錯誤的空間特征。然后,在對特定的下游任務進行微調時,我們將卷積核重新填充為3*3大小。假設第i卷積層的預訓練中的權重為(忽略通道維),填充內核如下
其中是MAE預訓練學習到的值,初始化為0。此外,我們在ViTAE模型中采用一種淺層PCM的設計,其依次為卷積層、批歸一化層、SiLU層和卷積層,以節省內存占用。圖3顯示了用于MAE預訓練的ViT和ViTAE網絡中的基本模塊。
圖3:MAE編碼器中采用的塊結構
(a) ViT的基本模塊,(b) 改進后的ViTAE Normal Cell
我們使用“Base”版本的ViT和ViTAE,它們都具有約一億參數。這兩種網絡被分別表示為“ViT-B”和“ViTAE-B”。其詳細結構見表1,其中“Patch Size”表示特征嵌入層的圖像塊尺寸,“Embedding Dim”表示Token的維度,“Head”表示MHSA中SA的個數,“Group”表示PCM中分組卷積的組數,“Ratio”指FFN的特征維膨脹率, “Depth”表示兩種網絡中堆積模塊的數量。
表1 :ViT-B和ViTAE-B的超參數設置
2.4 采用RVSA進行微調
與自然圖像相比,遙感圖像通常尺寸更大。由于完全自注意力具有和圖片分辨率呈平方關系的計算復雜度,直接將采用完全自注意力的預訓練模型應用于下游任務時會顯著增加訓練成本。為此,我們在微調階段采用窗口自注意力替換原始的完全自注意力,這將計算代價降低到與圖像大小線性相關的復雜度。因為這種替換只改變了參與自注意力計算的Token范圍,而不引入新的參數,因此可以在預訓練-微調范式中直接轉換。然而,原始的窗口自注意力在固定水平和垂直方向上采用固定大小的窗口,這與遙感圖像中目標的任意朝向特點不符,導致使用固定方向固定大小的窗口可能并非最優。為此,我們設計了RVSA。
具體來說,我們引入了一系列變換參數來學習可變方向、大小和位置的窗口,包括相對參考窗口的偏移量、尺度縮放因子以及旋轉角度。具體地,給定輸入特征,首先將其劃分為幾個不重疊的參考窗口,即每個窗口的特征為(其中表示窗口大小),總共得到個窗口。然后,我們通過三個線性層去獲得查詢特征,初始的鍵特征和值特征。我們用去預測目標窗口在水平和豎直方向上的偏移和縮放,以及旋轉角度
GAP是全局平均池化操作的縮寫。以窗口的角點為例
上式中,表示初始窗口左上角和右下角的坐標,表示窗口的中心坐標,分別是角點與中心在水平和垂直方向上的距離。我們用估計到的參數來對窗口進行變換,
是變換后窗口的角點坐標。然后,從變換后的窗口中采樣鍵特征,從而和查詢特征一起計算自注意力。采樣的鍵特征和值特征中Token的數量與查詢特征中Token的數量相同,從而保證RVSA與原始窗口自注意力機制具有相同的計算復雜度。
這里,是一個窗口中一個SA的輸出特征,,是SA的個數。然后,沿著通道維度連接來自不同SA的特征,并且沿著空間維度連接來自不同窗口的特征,以恢復輸入特征的形狀,最終獲得RVSA的輸出特征,圖4展示了RVSA的示意圖。
圖4:第個SA上所提出RVSA的完整流程
RVSA的變體:我們還提出了RVSA的一種變體,此時鍵特征和值特征可以來自不同窗口,即我們分別預測鍵特征和值特征窗口的偏移,縮放和旋轉因子
這個更靈活的架構被稱為RVSA。
為了使MAE預訓練模型適應遙感下游任務,我們將原始Plain ViT中的MHSA模塊替換為RVSA。按照ViTDet [10]中的策略,我們在每1/4個深度層采用完全自注意力。由于ViT-B 和 ViTAE-B有12 層,因此我們在第3、6、9和12層使用完全自注意力,并在所有其他層采用RVSA。修改后的網絡分別表示為“ViT-B + RVSA”和“ViTAE-B + RVSA”。圖5展示了 ViT-B + RVSA和ViTAE-B + RVSA中替換注意力后模塊的結構。我們也對比了采用普通窗口自注意力、VSA和RVSA的變體。它們被分別表示為“ViT-B-Win”、“ViT-B + VSA”、 “ViT-B + RVSA”、“ViTAE-B-Win”、“ViTAE-B + VSA”和“ViTAE-B + RVSA”。
圖5:修改注意力后模塊的結構(a)ViT-B+RVSA。(b)ViTAE-B+RVSA
最后,我們在圖6中展示了上述預訓練和微調過程的完整框架,以便于讀者理解所提出的方法。
圖6:預訓練和微調的流程(以RVSA為例)
03
實驗結果
我們將所提出的模型在包括場景分類、對象檢測和語義分割等多個遙感任務上進行實驗,并且還進一步展示了其在計算復雜度、遷移學習的數據效率以及可解釋性等方面的優勢。
3.1 目標檢測
表2:不同先進方法在DOTA-V1.0數據集上的精度。
表3:不同先進方法在DIOR-R數據集上的精度
我們將所提出的方法與迄今為止最先進的一些方法進行了比較,結果列于表2和表3中。每列精度前三分別用粗體,紅色和藍色標記。在DOTA-V1.0數據集上,我們分別列出了單尺度訓練和多尺度訓練的結果。在單尺度訓練設置上,我們的模型在五個類中表現最好,超過了以前的最佳方法約1%的mAP。在競爭更激烈的多尺度訓練中,我們的模型在總共四個類別中獲得第一。特別的,我們的方法在一些具有挑戰性的類別(如環島和港口)中的檢測結果顯著優于之前的方法,從而在DOTA-V1.0上取得了新的精度記錄,即81.24%的mAP。在更具挑戰性的DIOR-R數據集上,我們的模型在11個類別中表現最好。與現有方法相比,其檢測性能提高了10%以上,并以5% mAP的優勢顯著超過第二名。值得注意的是,我們成功地證明了建立強大的Plain ViT基線的可能性:事實上,ViT-B+VSA和ViTAE-B+VSA在DOTA-V.1.0和DIOR-R數據集上已經超過了之前的方法并取得了很好的檢測性能。當進一步引入旋轉機制后,它們的性能仍然能被進一步提高。
3.2 場景分類
表4:不同方法在場景分類任務上的精度
表4展示了不同模型在場景分類任務上的結果。在此任務中,MAE 預訓練的ViT-B在大多數設置上獲得最佳效果,因為所有Token都參加了MHSA計算 ,這種方式提取的全局上下信息有利于場景識別。我們的 RVSA 模型在三個設置(包括 UCM-55、AID-28 和 NWPU-28)中優于以前的方法。而在其他設置中,我們的模型可以與探索研究院先前提出的當前最先進的模型:即在 MillionAID上采用有監督預訓練的層次化模型RSP-ViTAEv2-S [5]相媲美。與VSA 方法相比,我們所提出的模型主要在NWPU-19設置中表現較差。這是因為相比VSA,RVSA 需要相對更多的訓練數據來學習最佳窗口配置,而NWPU-19 的訓練數據規模相對較小。當采用較大規模數據集,如NWPU-28 的設置時,我們的模型超越了ViT-B + VSA,ViTAE-B + VSA和RSP-ViTAEv2-S等先進模型。
3.3 語義分割
表5:不同方法在語義分割任務上的精度
表5顯示了不同分割方法的結果。我們的模型獲得了與當前最佳方法相當的性能。盡管如此, 我們也必須承認其在分割任務上的性能不如在檢測和場景分類任務上令人印象深刻。我們認為這有兩個原因。首先,我們使用經典但簡單的分割框架 UperNet,它不能有效地將高級語義信息傳播到高分辨率特征圖上。另一個原因是我們采用的視覺Transformer 骨干網絡直接通過的圖像塊來嵌入編碼網絡特征,并且特征圖分辨率始終保持輸入大小的1/16,這可能會丟失細節,不利于像素級語義分割任務。盡管如此,我們提出的RVSA仍然可以提升Plain ViT的性能并達到與層次化模型RSP-ViTAEv2-S相當的性能,且優于ViT-B、ViTAE-B 和 VSA等模型,證明了其從可變窗口中學習有用上下文信息的強大能力。
3.4 計算復雜度
表6:不同模型的復雜度和訓練代價
我們以ViT為例,在DIOR-R檢測數據集上比較了不同模型的復雜度和訓練代價。表6列出了包括參數數量 (Params)、計算量 (FLOPs)、GPU 內存在內的多種評估指標,所有模型參數量均超過1億。由于完全自注意力的二次復雜度,ViT-B具有最大的內存占用,最大的FLOPs以及最長的訓練時間,因此需要使用兩個GPU才能在相當的時間完成訓練。ViT-B-Win通過采用窗口自注意力緩解了這些問題。需要注意的是,ViT-B + VSA的FLOP比ViT-B-Win小,這是因為填充(padding)操作是在生成查詢特征、鍵特征和值特征之后實現的。由于可學習的縮放和偏移因子,ViT-B + VSA比ViT-B-Win略多一些內存占用。與ViT-B+VSA相比,ViT-B+RVSA具有相似的復雜度,而ViT-B+RVSA略微增加了參數和計算開銷,因為它對鍵特征和值特征分別預測窗口。與ViT-B相比,所提出的ViT-B + RVSA和ViT-B + RVSA可以節省大約一半的內存并加快訓練速度,同時還具有更好的性能。
3.5 遷移學習的訓練數據效率
圖7:不同的模型在DIOR-R數據集上采用不同比例的訓練集進行訓練的結果
遷移學習的訓練數據效率是衡量基礎模型的一項重要能力。在這里,我們在DIOR-R 數據集上,使用不同數量的訓練數據進行實驗。我們通過分別隨機選擇原始訓練集的 20%、40%、60% 和 80% 的圖像來獲得一系列較小的訓練集。然后,我們分別在這些數據集上微調預訓練模型,并在原始測試集上對其進行評估。為了便于比較,我們也訓練了一些小規模模型,例如 RSP-ResNet-50、RSP-Swin-T和RSP-ViTAEv2-S,它們采用訓練集中所有的數據進行訓練。圖7顯示了相關結果。可以看出,無論訓練樣本的數量如何,所提出的模型都優于相應的ViT-B和ViTAE-B基線模型。由于我們考慮了遙感圖像中任意方向的對象,所提出的具有可學習旋轉機制的RVSA在大多數情況下都可以超越VSA。此外,它們僅使用40%的訓練樣本就達到了與Swin-T相當的性能,當使用60%的訓練樣本時,它們的性能優于ResNet-50和Swin-T。當采用80%的訓練樣本時,它們超過了強大的骨干網絡ViTAEv2-S。上述結果表明我們的模型在遷移學習時具有良好的訓練數據效率。
3.6 窗口可視化
圖8:不同注意力方法生成窗口的可視化。(a) 窗口自注意力。(b) VSA。(c) RVSA。(d)和(e)分別是RVSA為鍵特征和值特征生成的窗口
以ViT為例,圖8 可視化了不同的網絡在倒數第二個模塊中的注意力層生成的窗口。可以看出,VSA生成的窗口可以縮放和移動以匹配不同的對象。然而,VSA 無法有效處理遙感圖像中任意方向的目標,例如圖8第二行中傾斜的飛機。相比之下,我們的RVSA引入了旋轉因子來解決這個問題,獲得更多樣化的窗口并有利于提取更豐富的上下文信息。同樣值得注意的是,每個頭可以產生不同方向和位置的窗口,并來覆蓋特定角度和位置的飛機。因此,通過使用多頭注意力的方式,圖片中不同方向的飛機可以被不同頭的窗口覆蓋,這意味著RVSA可以更好地處理任意方向的物體。與RVSA相比, RVSA進一步提高了生成窗口的靈活性。通過將 (d) 和 (e) 與 (c) 進行比較,我們可以發現鍵特征和值特征的窗口形狀略有變化,這在擁有大量可用的訓練數據和處理具有挑戰性的樣本時比較有用。通過將學習到的窗口進行可視化,我們提供了一種分析所提出的模型的工作機制的手段,可以增強其學習過程和學習結果的可解釋性。
04
總結
本工作提出了全球首個面向遙感任務設計的億級視覺Transformer大模型。具體來說,我們首先基于具有代表性的無監督掩碼圖像建模方法MAE對網絡進行預訓練來研究Plain ViT作為基礎模型的潛力。我們提出了一種新穎的旋轉可變大小窗口注意力方法來提高Plain ViT的性能。它可以生成具有不同角度、大小、形狀和位置的窗口,以適應遙感圖像中任意方向、任意大小的目標,并能夠從生成的窗口中提取豐富的上下文信息,從而學習到更好的物體表征。我們在典型的遙感任務上對所提出的模型進行實驗,結果證明了Plain ViT作為遙感基礎模型方面的優越性和有效性。我們希望這項研究可以為社區提供有價值的見解,并激發未來對開發遙感基礎模型的探索,尤其是基于Plain ViT的研究。
審核編輯 :李倩
-
視覺
+關注
關注
1文章
147瀏覽量
23937 -
數據集
+關注
關注
4文章
1208瀏覽量
24691 -
Transformer
+關注
關注
0文章
143瀏覽量
5997 -
大模型
+關注
關注
2文章
2427瀏覽量
2649
原文標題:全球首個面向遙感任務設計的億級視覺Transformer大模型
文章出處:【微信號:CVer,微信公眾號:CVer】歡迎添加關注!文章轉載請注明出處。
發布評論請先 登錄
相關推薦
評論