基于改進Swin Transformer的艦船目標實例分割算法

錢 坤, 李晨瑄, 陳美杉, 郭繼偉, 潘 磊

(1. 海軍航空大學岸防兵學院, 山東 煙臺 264000; 2. 中國人民解放軍32127部隊, 遼寧 大連 116100; 3. 航天工程大學研究生院, 北京 101416)

0 引 言

近年來,隨著計算機圖形算力的爆發式增長和人工智能算法的長足進步,計算機視覺領域進入了新的技術時代,圖像分割作為計算機視覺的重要分支,是根據一定的規則將圖像中的像素分成不同的部分,并添加相應標簽,可以視為是圖像分類識別向像素級的延伸,而其中的實例分割兼具了目標檢測和語義分割的雙重任務,在實現上更具挑戰意義和實用價值。目前,實例分割已經廣泛應用于汽車自動駕駛、智能輔助醫療和遙感影像判讀等民用領域。在軍事應用上,實例分割也是精確制導、偵察監視等研究領域的核心問題,對于推進軍事現代化和信息化有著重要意義[1]。

早期的實例分割受益于目標檢測相關研究成果的延伸,Hariharan等首次使用同時檢測和分割算法[2]通過生成建議框、特征提取、區域分類與改良等步驟同時完成了檢測和分割操作。He等人提出的掩模區域卷積神經網絡(mask region-convolutional neural network, Mask R-CNN)[3]是一種功能強大的基線算法,通過添加預測興趣區域的語義分割分支能夠高效地完成實例分割任務,也引領了后續一系列改進算法的產生,如掩模評分區域卷積神經網絡(mask scoring R-CNN, MS R-CNN)[4]、邊界保持Mask R-CNN(boundary-preserving Mask R-CNN, BMask R-CNN)[5]等。Long等人提出的全卷積網絡(fully convolutional networks, FCN)[6]是一種端到端的圖像分割方法,通過使用全卷積網絡、反卷積層和跳躍連接實現了圖像分割,成為該領域的代表性算法。為進一步提升FCN算法的精度,Bolya等人提出了只看系數算法(you only look at coefficients, YOLACT)[7],該算法通過兩個并行的子網絡來實現實例分割,一個用來得到類似FCN的分割原型圖,另一個用來得到檢測框,然后融合原型圖和檢測框后得到掩碼,其在處理速度上實現了突破,成為首個能夠完成實時實例分割的算法,隨后的YOLACT++算法[8]繼續優化了算法結構和執行效率,進一步提升了算法速度和精度。

近兩年來,發軔于自然語言處理(natural language processing, NLP)的Transformer[9]模型在計算視覺領域大放異彩,在實例分割研究中相較傳統CNN展現出了極強的競爭力。基于Transformer的實例分割算法[10]是首個基于Transformer的實例分割框架,通過使用循環細化策略進行檢測和分割,為實例分割提供了新角度。隨后的基于Transformer的圖像分割算法[11],使用Transformer預測每個實例類別,并動態生成具有多個分割的掩碼級上采樣模塊,該算法與原始Transformer相比在運算速度和資源開銷上更具優勢。基于滑動窗口的Transformer(shifted windows Transformer,Swin Transformer)[12]實例分割算法,使用滑動窗口策略和層級化設計,其將注意力計算限制在一個窗口中,旨在引入與CNN卷積類似的局部性操作,并顯著降低計算量,進一步提高了實例分割的速度和精度。隨著對Swin Transformer模型研究的不斷深入,如何在保持較低運算開銷條件下盡量充分地實現上下文信息的交互,如何理解算法中自注意力機制的作用方式以及如何對其他模塊進行優化等問題的提出,為進一步提升Swin Transformer模型性能提供了可能的方向。

海戰場艦船目標精確識別與分割是反艦導彈末段圖像制導的核心問題,對于艦船目標的檢測識別算法很多,但細化至像素級的實例分割研究相對較少,已有研究也多是針對合成孔徑雷達(synthetic aperture radar,SAR)圖像進行實例分割[13],而SAR圖像多為高空俯視視角,從俯視視角得到的艦船輪廓相對單一,對于數據集的準備和訓練也相對容易。目前,現役主戰反艦導彈為保證隱蔽突防的成功率,多采用低飛掠海攻擊模式,從艦船側視角接近目標時,艦船目標的尺度和角度變化更加劇烈,因此對于算法的要求也更高。為進一步提升海戰場目標檢測能力,實現精確化打擊,對于將一般的目標檢測識別延伸細化至像素級精度的實例分割,有著很強的實戰意義。

本文提出了一種將改進Swin Transformer作為骨干網絡的艦船目標實例分割算法,在Swin Transformer模塊前增加了局部增強感知模塊,旨在提高網絡的上下文信息交互能力并拓展感受野,提高分割精度;在特征融合中使用反向特征金字塔(reverse feature pyramid,RevFP)網絡,在優化分割準確率的同時進一步提升圖像處理速度;使用在線困難樣例挖掘方法,改善自制數據集難易樣本失衡問題,優化算法在小樣本數據集條件下的訓練效果;在經過自制的艦船目標數據集訓練后,相關實驗結果表明,算法能夠實現多角度、多尺度艦船目標實例分割,分割邊緣清晰準確,相較基線算法兼具精度和速度優勢。

1 基于Swin Transformer的實例分割算法

1.1 Transformer模型結構

Transformer模型最先在NLP任務中被廣泛應用,在NLP任務中需要編碼器抽取很多特征。其中,首先需要考慮的就是每個詞的上下文語義,因為每個詞的具體含義都跟上下文強相關,而上下文又分為方向和距離兩個屬性,傳統的循環神經網絡(recurrent neural network,RNN)[14]只能對句子進行單向編碼,CNN只能對短句進行編碼,而Transformer既可以同時編碼雙向語義,又能夠抽取長距離特征,所以在上下文語義抽取方面要優于RNN和CNN。NLP任務需要抽取的第二種特征是序列的順序,在這方面Transformer的表現一般介于RNN和CNN之間。NLP任務對計算速度是比較敏感的,由于RNN無法并行處理序列信息,因此表現最差,CNN和Transformer都可以進行并行計算,但Transformer模型結構稍顯復雜,所以在速度方面稍遜CNN。綜上,由于Transformer在效果和速度方面性能表現均衡,所以在NLP任務中很快便脫穎而出,之后隨著研究的深入,Transformer被引入其他任務中,均有不俗表現,已隱隱呈現出成為高效的通用計算架構的趨勢。

Transformer模型是一種典型的編碼器-解碼器結構,其中最為重要是多頭自注意力(multi-head self-attention,MSA),殘差連接和歸一化(add &layer normalization,Add &LN)以及前饋網絡3個模塊。MSA負責將輸入投影到不同空間,得到Q、K、V這3個矩陣,投影操作是避免后續矩陣Q和V直接進行點積操作,并使輸入的每個部分能夠融合上下文的語義,增加多樣性,之后再進行注意力的計算。在Add &LN模塊中先對注意力計算的輸出進行投影和隨機舍棄運算,加上原始輸入后再進行LN操作,添加原始輸入是參考了殘差連接的思想,相當于在求導時增加了一個恒等項,達到避免發生梯度消失的目的。同時,歸一化的做法在神經網絡中主要起到提升網絡泛化性的作用,另外,將歸一化設置在激活函數之前,也能夠有效避免數據落入飽和區,進一步減少梯度消失的情況。前饋網絡模塊主要作用是給Transformer提供非線性變換,提升擬合能力,在這一過程中選擇了高斯誤差線性單元(Gaussian error linear unit, GeLU)函數[15]作為激活函數,因為GeLU函數引入了正則思想,越小的值越有可能被舍棄,相當于線性修正單元函數(rectified linear units, ReLU)[16]和隨機舍棄的綜合,因為ReLU函數的值只有0和1,所以單純使用ReLU就缺乏這樣的隨機性。

1.2 基于Swin Transformer的實例分割算法原理

Swin Transformer的創新性思想主要表現在層次性、局部性以及平移不變性。層次性體現在特征提取階段使用了類似CNN中的層次化構建方法,對輸入圖像分別進行4倍、8倍和16倍的下采樣,得到多尺度特征圖。局部性主要體現在自注意力計算過程中,將計算約束在劃分好的局部非重疊窗口內進行,使得算法復雜度從此前的和圖像大小的平方關系變成了線性關系,大幅降低了計算量,提升了算法效率。在前后兩層的Transformer模塊中,非重疊窗口的劃分相比前一層做了半個窗口的偏移,這樣使得上下兩個窗口的信息得到了有效交互,相比于CNN中常見的滑動窗口設計,保持了平移不變性的同時也沒有造成準確率的降低。

Swin Transformer的算法流程大致如下,輸入圖像為H×W大小的RGB(red, gree, blue)三通道圖像,首先通過圖塊分割模塊進行分塊操作,分為4×4大小的圖塊,之后沿3個通道方向進行展平,展平后大小為4×4×3=48,所以通過圖塊分割后圖像由[H,W,3]變為[H/4,W/4,48],而后通過線性嵌入模塊對每個像素通道數據進行線性變換,得到[H/4,W/4,C],以上操作均通過二維卷積實現,具體結構如圖1所示。

圖1 Swin Transformer網絡結構Fig.1 Architecture of Swin Transformer net

隨后數據進入Swin Transformer模塊,經歸一化后,進入窗口多頭自注意力(windows multi-head self-attention,W-MSA)計算模塊,將圖像塊進一步劃分成互不重疊的區域,并在區域內計算自注意力;多層感知機(multilayer perceptron,MLP)是一個使用GeLU函數作為激活函數的前饋網絡,旨在完成非線性變換,提升算法擬合能力。由于在W-MSA中,只會在每個窗口內進行自注意力計算,所以窗口與窗口之間是無法進行信息傳遞的。為此,引入了偏移(shifted W-MSA, SW-MSA)模塊,在第L層劃分互不重疊窗口后,在第L+1層內偏移半個窗口距離重新劃分窗口,使得不同層中有部分窗口的信息得以交互,以上操作可以歸結為第一階段,Swin Transformer模塊結構如圖2所示。

圖2 Swin Transformer模塊結構Fig.2 Architecture of Swin Transformer blocks

為了得到多尺度的特征信息,還需要構建一個層級式結構,Swin Transformer網絡中設計了圖塊拼接模塊,來實現類似CNN中池化操作,經過拼接過后圖像的寬高尺寸減半,通道數翻倍,而后再通過Swin Transformer模塊,這樣的結構以串聯的形式相連,即圖1中的第2階段到第4階段。

為應對不同需求,Swin Transformer實現了網絡結構的系列化,其網絡由簡單到復雜依次是Swin-Ting,Swin-Small,Swin-Baes和Swin-Large,網絡越復雜性能越強,但相應的計算開銷也越大,圖1就是以Swin-Ting為例繪制的。

2 基于改進Swin Transformer的艦船目標實例分割算法

2.1 算法框架

本算法使用Swin Transformer作為骨干網絡,完成特征提取工作,圖像進入圖像分割模塊后形成序列化小塊,并沿通道方向展平,通過線性嵌入模塊完成線性到非線性再到線性的3次變換;在進入Swin Transformer塊前,先通過局部增強感知(local enhanced sensing,LES)模塊,旨在進一步抑制無效特征,增強有效參數,提高上下文信息之間的交互;在特征融合階段使用RevFP網絡[17],在采用局部融合操作的同時,集成高級與低級特征,獲得更好的融合效果和更低的延遲;在網絡的訓練階段,使用在線困難樣例挖掘(online hard example mining, OHEM)[18]方法,解決自制數據集難易樣本不均衡,訓練效果不理想的問題,通過設置負樣本池反復迭代,使算法在小樣本數據集上獲得更好的效果;最后使用基于Mask R-CNN的檢測頭,完成實例分割任務。算法整體框架如圖3所示。

圖3 算法框架Fig.3 Algorithm framework

2.2 局部增強感知模塊

在Swin Transformer模型中,為了在與圖像大小成線性關系的計算復雜度下盡可能地實現上下文語義信息的“全局”范圍交互,在層次化結構中使用偏移窗口方案來實現不同窗口間的信息交互。但實際上,每次偏移半個窗格,在空間上依然存在大量上下文信息不能很好地交流,為此,本文提出了在Swin Transformer模塊前插入一個局部增強感知模塊,旨在進一步提升信息交互能力,拓展感受野尺寸,具體結構如圖4所示。

圖4 帶有LES模塊的Swin Transformer骨干網絡結構Fig.4 Swin Transformer backbone network structure with LES block

在卷積神經網絡中,卷積核大小決定了卷積感受野的尺寸,而與之對應的是不同尺寸的感受野適合識別分割不同尺寸的目標,由于反艦導彈在接近艦船目標過程中,目標變化的過程具有多角度、多尺度特性,所以設計能夠融合多尺度的感受野對于提升識別和分割精度有著重要作用。受“深度實驗室”系列算法[19-22]啟發,設計了包含3條并行支路的局部增強感知模塊。

在Swin Transformer中數據流由向量構成,數據首先進行特征向量的重新投影,形成多維空間特征映射,然后通過并行的空洞卷積(dilated convolution, DConv)[23]分支,其中,每個支路中核大小為1×1的普通卷積旨在降低通道數,3條分支的空洞卷積擴張率分別為r=1,r=3和r=5,根據空洞卷積卷積核計算公式K=(r-1)(k-1)+k。其中,K為空洞卷積卷積核尺寸,r為擴張率,k為常規卷積卷積核尺寸,由此分別獲得3×3,7×7和11×11大小的感受野,利用DConv來增加感受野尺寸并不以犧牲特征分辨率為代價,且可以在不同的尺度上對大范圍的上下文信息進行良好編碼,使特征圖獲取更精確的語義與定位信息,有效增強尺度、角度劇烈變化時的艦船目標識別分割的魯棒性[24]。

數據通過3條擴張率不同的空洞卷積分支,X為輸入特征,Y為輸出特征,3條支路的輸出如下所示:

(1)

式中:Conv為卷積核為k的普通卷積運算;DConv為擴張率為r的DConv運算。

通過并行的3條支路后,輸出特征經通道拼接,通過GeLU激活函數,并采用殘差網絡思想,增加跳躍連接,隨后輸出至Swin Transformer模塊。

2.3 基于RevFP網絡的特征融合

在特征融合階段,特征金字塔網絡[25](feature pyramid networks,FPN)是目前最為常用的選擇,其結構如圖5(a)所示。后續相關研究提出了多種不同的架構,如路徑聚合網絡(path-aggregation network,PANet)[26],雙向FPN(bi-directional FPN,BiFPN)[27]等,這些方法使用了不同形式的雙向特征融合。相關實驗表明,以上改進結構均能有效改善網絡性能,已在目標檢測、實例分割領域得到了廣泛應用,但這些網絡通常是以固定順序堆疊特征金字塔,更長的信息鏈路,會降低推理速度。此外,由于FPN中僅在近鄰層級特征進行了局部融合,這就導致非相鄰特征的語義信息會被稀釋。文獻[17]提出了RevFP網絡,這是一種利用局部雙向特征融合的簡化的雙向金字塔推理架構,相比于其他雙向金字塔架構,RevFP具有更好的性能、更小巧的模型和更低的延遲,RevFP結構如圖5(b)所示。

圖5 FPN和RevFP特征融合網絡結構Fig.5 FPN and RevFP feature fusion network structure

傳統的FPN網絡采用自頂向下的方式將高層特征傳遞至低層特征,底層特征包含豐富的幾何性狀信息,而高層特征則包含更多的語義信息,FPN網絡的信息傳遞可以描述為如下形式:

Pi=fi(Ci,Pi+1),i∈{2,3,…,6}

(2)

為進一步促進高效的雙向特征融合,在RevFP網絡中采用局部融合操作同時集成高級與低級特征:

Pi=fi(Ci,Ci+1,Pi-1)

(3)

式中:Ci為特征提取骨干提取的特征,i為階段數,即圖像特征分辨率減半的次數。如C3表示特征提取第3階段,輸出為原圖像1/8尺寸的特征圖,對應融合后的輸出記作Pi。

2.4 在線困難樣例挖掘

在線困難樣例挖掘主要用于解決數據樣本不均衡或簡單樣本過多的問題。通常情況下,在訓練網絡的最后階段,當前的網絡模型會對所有訓練集中的樣本進行測試,逐個分析其當前的損失函數值,并將樣本按照置信概率誤差進行降序排列,其中誤差較大的樣本就是所謂的困難樣例。

在理想狀態下,一個數據集的難易樣例設置應該是相對平衡的,但在小規模或自制數據集中,樣本難易不平衡的現象幾乎普遍存在。比如掠海攻擊視角,艦船目標角度變化明顯,導致相同部位在不同角度上觀察變化巨大;再如,數據集規模較小時,對于某類目標的特征學習不夠,會導致損失函數和訓練效果退化。為提高數據集利用效率和算法精度,受數據增強和在線困難樣例挖掘思想啟發,設置了負樣本檢測池,在訓練20輪后,將每批次檢測準確率小于0.5的樣本置于負樣本檢測池中,累計64個負樣本后,將其重新投入網絡進行訓練,反復迭代以改進算法對于困難樣本的特征提取能力。

3 實驗結果

3.1 實驗環境及數據集

為模擬反艦武器掠海攻擊模式末段圖像制導呈現的影像,選用文獻[28]提出的基于水平或低視角的艦船目標數據集,數據集共有1 554幅圖像,作戰艦艇區分航空母艦和驅逐艦,分別有636幅和741幅,民用船只包含漁船、液化天然氣船、集裝箱船、快艇、帆船等類別,共558幅,由于網絡支持多分辨率圖像訓練,所以沒有統一圖像大小,所有圖像大小介于500×680至1 280×960之間。

為強化反艦武器打擊效果,標注過程中將艦體主體作為重點標注目標,艦體外緣突出的垂直天線、艦炮炮管等細長狀目標,以及甲板欄桿、雷達天線等存在較大間隙的目標不作為艦船主體進行標注,標注示例如圖6所示。

圖6 標注示例Fig.6 Annotation example

經統計,共標注1 052幅圖像,每圖像平均標注1.23個目標。標注使用開源工具labelme,并將標注后生成的.json格式的文件使用腳本轉換成微軟通用目標語義(Microsoft common objects in context, MS COCO)數據集[29]標準格式,并按照MS COCO數據集的目錄構成要求進行文件的組織和保存。另外,在訓練前還需要將標注好的數據集按照8∶1∶1的比例隨機劃分為訓練集、測試集和驗證集,3個子集的數量分別為842幅、106幅和104幅。具體標注情況如表1所示。

表1 數據集標注信息統計表Table 1 Statistical table of dataset annotation information

算法依托深度學習主流開源框架PyTorch,使用NVIDIA CUDA GPU加速神經網絡的訓練和推斷,具體實驗環境如表2所示。

表2 實驗環境Table 2 Experimental environment

3.2 評價指標

在圖像實例分割研究中,常用MS COCO數據集中的一些評價指標來衡量算法的優劣,最主要的指標是平均精度(average precision,AP),表示對某一類目標檢測精度,AP由交并比(intersection-over union,IoU)定義,IoU表示兩個實例掩碼的重疊度。其他比較重要的性能指標還有參數量和幀率(frame per second, FPS),參數量由網絡模型結構確定,單位為M,FPS主要用來衡量算法執行的速度,常用的評價指標如表3所示。

表3 MS COCO數據集主要評價指標Table 3 Main evaluation indexes of MS COCO dataset

3.3 實驗結果與分析

將網絡在艦船目標數據集上進行訓練,待網絡收斂后,得到改進后的網絡模型和相應的權重文件,為驗證算法性能,在相同實驗平臺和訓練數據條件下,選擇5個經典的實例分割算法進行性能對比,分別是FCN、Mask R-CNN、級聯Mask R-CNN(Cascade Mask R-CNN)、YOLACT++算法和以Swin-Ting為骨干網絡的Swin Transformer算法。由于本文算法的骨干網絡是在Swin-Ting基礎上改進而來,所以在實驗對比對象的選擇上也盡可能選擇網絡模型大小相近的算法進行對比,具體實驗結果如表4所示。

表4 算法對比結果Table 4 Algorithm comparison results

圖7為本文算法與基線算法的實例分割結果對比圖,圖7(a)組中放大區域顯示本文算法相較基線算法對目標邊緣的分割效果更加準確;圖7(b)組中基線算法出現了誤檢;圖7(c)組中兩個算法均準確分割出目標,但本文算法給出的目標置信概率顯著高于基線算法;圖7(d)組基線算法出現了錯檢,將航空母艦同時識別成了航空母艦和驅逐艦。

圖7 對比實驗結果示例Fig.7 Examples of comparative experimental results

3.4 消融實驗

為驗證不同模塊對于實例分割算法的影響,設計了消融實驗,消融實驗依然以原始Swin-Ting+FPN,且無OHEM輔助訓練的算法作為基線算法進行對比。6組消融實驗是單一模塊或兩兩模塊的組合算法,最后一組是本文的完整算法,即改進Swin-Ting+RevFP+OHEM,具體實驗結果如表5所示。其中,“×”表示未使用該模塊;“√”表示使用了該模塊進行改進優化。

表5 消融實驗數據Table 5 Ablation experiment data

實驗結果表明,單一或組合模塊的改進在分割準確度或處理速度上大多可以產生正向優化效果,證明對原始基線算法的改進是有效的。其中,LES模塊對于網絡分割準確度的提升比較明顯,但也會相應增加一定計算開銷,降低圖像處理速度;RevFP模塊的主要貢獻在于對算法處理速度的提升,對準確度的提升相對有限;在基線算法基礎上單獨應用OHEM方法即可對算法精度有一定提升,證明了數據集的樣本難易平衡性能夠直接影響網絡訓練效果,在規模較小的數據集上能夠取得相對優異的訓練效果,這對于樣本容量和質量比較有限的自制數據集有著重要意義。為進一步驗證OHEM方法在該方面的作用,還進行了60%和80%訓練集的實驗,結果如表6所示。

表6 不同規模訓練集訓練效果對照實驗結果Table 6 Comparative experimental results of training effects of different scale training sets %

實驗結果表明,在數據集樣本容量相同情況下,使用了OHEM方法進行訓練的Swin-T網絡訓練結果均優于基線算法,在80%數據集條件下,mAP高出0.9%,在60%數據集條件下,mAP高出1.2%,在數據集容量有限的情況下,優勢有擴大趨勢。

4 結束語

傳統艦船目標識別多使用檢測框標示出目標,為將目標識別延伸至像素級別,本文提出了一種基于改進Swin Transformer的艦船目標實例分割算法,可以實現多尺度,多角度,基于掠海攻擊視角的多類別艦船目標實例分割,算法分割精度、速度均優于基線算法,結合消融實驗,論證了對算法改進的有效性。

(1) 針對傳統Swin Transformer模型存在圖像上下文信息交互不充分的問題,設計了由并行DConv分支組成的局部增強模塊,拓展感受野尺寸和語義交互能力,使算法mAP提高0.9%。

(2) 針對傳統FPN特征融合網絡存在非相鄰的特征語義信息被稀釋的問題,使用RevFP網絡進行替代,在分割精度和速度上均有不同程度提升。

(3) 針對自制數據集樣本數量較少,樣本難易程度不平衡導致的訓練結果不佳的問題,使用OHEM方法反復迭代以改善訓練效果,切實提升了小樣本數據集的訓練效果。

下一階段的研究重點是進一步提高網絡模型的特征提取能力,并建立樣本數量更大的艦船目標數據集,重點解決復雜島岸背景下的艦船目標識別或實例分割問題。同時,將算法遷移至嵌入式平臺端,對彈載平臺的分割效果進行實驗驗證。