基于深度學習的Attention U-Net語義分割模型研究

薛澤民,鄒連旭,黃志威,冉杰,余若巖,鄭國勛

(1.長春工程學院,長春 130012; 2.長白山歷史文化與VR技術重構吉林省重點實驗室,長春 130012)

0 引言

近年來,深度學習在計算機視覺領域大放異彩,已逐漸成為主流研究方向,越來越多的研究人員將其深度學習技術應用于圖像語義分割。在深度學習進行圖像處理的研究中,通常是將圖像輸入到編碼器中,得到1個中間的上下文,再使用解碼器對該上下文進行解碼,最后還原成1個輸出的序列。而目前的做法是使用1個循環神經網絡實現編碼器,用另外1個循環神經網絡去實現解碼器。針對編碼器和解碼器的實現,還有一些人采用不同模型結合的方法,比如使用雙向循環神經網絡或者采用帶有長短期記憶的雙向循環神經網絡等[1]。這些方法均需要將大量的信息壓縮成1個固定長度的中間上下文向量,因此使得大量的信息在被壓縮后產生損失,尤其是先輸入的信息會被后輸入的信息稀釋,且隨著輸入的序列越長,問題越嚴重,最后會導致解碼質量嚴重下降。

Jay Alammar為解決這種編碼器-解碼器結構存在的問題,在傳統的編碼器中引入了注意力機制。加入注意力機制后,編碼器需要將源圖像編碼成一個個向量。與傳統的編碼器-解碼器結構不同的是,傳統方式只需要把最后1個向量當作上下文向量傳給解碼器即可,但在注意力機制中,需要把所有的向量傳給解碼器,而在解碼時,會多一個額外的步驟,即通過注意力機制找出與當前輸出時間步長最相關的編碼器時間步長的輸入向量,再產生當前時間步長的輸出,使在產生每一個輸出的時候,都能夠做到充分利用輸入序列攜帶的信息,CHO等[2]采用這種方法在圖像語義分割中取得了非常不錯的成果。本文針對通過生成對抗網絡GAN擴充的數據集,提出了融入注意力機制的U-Net網絡訓練的圖像語義分割模型。

1 生成對抗網絡

生成對抗網絡是一種深度學習模型,被認為是近年來最有前景的無監督學習方法之一,適用于處理復雜分布的數據。與傳統方法需要標記大量的數據相比,生成對抗網絡可以在無需標注數據的情況下進行生成任務的學習。該網絡結構由生成器和判別器兩部分組成。首先生成器從潛在空間中進行隨機取樣并將其作為輸入,生成高度逼真的樣本。再把判別器接收的真實樣本和生成器生成的虛假樣本作為輸入,并努力將它們區分開。生成器和判別器相互對抗、不斷學習,模型通過這兩個模塊的博弈和優化逐步達到平衡點,從而實現擴充數據的目的[3]。為了保持生成器與判別器之間平衡的穩定性,即納什均衡,必須滿足式(1),結構如圖1所示。

圖1 生成對抗網絡結構圖

(1)

式中V(D,G)表示生成對抗網絡整體待優化的損失函數,z表示隨機變量,x表示真實數據,Pz表示生成數據的概率分布,P表示x的概率分布,D(x)表示判別x是否為真實的概率,D(G(z))表示判別重建數據是否真實的概率。

2 注意力機制

注意力機制是一種通過模擬人類視覺和認知系統的方法,通過注意力機制可以讓模型知道圖像中不同局部信息的重要性,從而提高模型的性能和泛化能力。主流的注意力機制采用編碼器-解碼器結構,如圖2所示,該結構將一個變長的輸入X=(x1,x2,…,xn)映射到一個變長輸出Y=(y1,y2,…,ym)。其中編碼器把一個變長的輸入序列X,通過非線性變換轉化為一個中間的語義表示C(C=f(x1,x2,…,xn))。解碼器則根據輸入序列X的中間語義表示C和先前已經生成的y1,y2,…,yi-1,預測并生成i時刻的輸出yi=g(y1,y2,…,yi-1,C)[4],在前述過程中,f()和g()為非線性轉換函數。

圖2 注意力機制的編碼器-解碼器結構

注意力機制包含多種,本文使用的是空間注意力?？臻g注意力是通過空間轉換模塊將原始圖片中的空間信息轉換到另一個空間,并保留關鍵信息,使關鍵區域的特征表達得到提升。它為每個位置生成權重掩膜并對特征進行加權輸出,從而增強感興趣的特定目標區域并減弱不相關的背景區域?？臻g注意力模塊的實現如圖3所示,首先對給定的H×W×C特征F′進行全局平均池化和全局最大池化操作,得到2個H×W×1的特征圖,分別代表不同信息;然后將這2個特征圖進行拼接,經過一個具有較大感受野的7×7卷積層進行特征融合;再通過Sigmoid操作生成一個權重圖Ms,并將其疊加到原始輸入特征F′上;最后根據權重圖對輸入特征進行縮放,從而完成增強目標區域的特征表達,計算過程如式(2)所示。

圖3 空間注意力實現圖

Ms(F)=σ(f7*7([AvgPool(F),MaxPool(F)])),

(2)

式中σ表示Sigmoid函數,f7*7表示7×7大小的卷積核,AvgPool表示平均池化,MaxPool表示最大池化。

3 融入空間注意力機制的U-Net

在U-Net網絡中,融入空間注意力模塊,可以使網絡更加關注感興趣的區域,并減少對不相關區域的依賴,我們稱其為Attention U-Net。即在U-Net的編碼器和解碼器之間插入空間注意力模塊,通過該模塊計算輸入特征圖的空間相關權重來調控信息的傳遞。具體來講,我們使用空間映射的方法對輸入特征圖進行卷積操作來獲取感興趣區域的權重分布,通過卷積層和激活函數來學習感興趣區域的空間特征,然后使用空間調整法將計算得到的注意力權重與輸入特征圖逐元素相乘得到經過調整的特征,最后將調整后的特征送入下一層網絡進行處理,從而增強感興趣區域的特征響應并有效抑制不相關區域的干擾。為了實現Attention U-Net,需要完成3部分內容:U-Net的實現、注意力門模塊的實現和U-Net與注意力門模塊的融合。

3.1 U-Net的實現

U-Net是一種基于FCN架構構建的U型網絡,其結構如圖4所示。左側部分是收縮路徑,經過一系列卷積和池化進行下采樣,提取圖像特征。右側部分是擴展路徑,進行上采樣,將特征圖與左側提取的特征相融合,在卷積層進行特征維度的調整,然后再進行上采樣和融合。重復前述過程直到上采樣到與原圖相同的尺寸時,添加1個卷積層用于對每個像素點進行分類來實現圖像的語義分割。這種結構充分利用了上采樣和下采樣過程中提取的多尺度特征信息以及跳躍連接的方式,能夠有效地獲取全局和局部的上下文語義信息,并將其應用于像素級的分割任務[5]。

圖4 U-Net結構圖

3.2 Attention U-Net的實現

3.2.1 實現注意力門模塊

圖5 注意力門模塊結構圖

(3)

式中σ1表示激活函數,*表示卷積操作,k表示卷積核參數,Fl表示在第l層的感受野內的所有位置的集合。

注意力系數α用來識別圖像中的顯著區域并修剪特征響應,它將輸入特征圖與注意力系數進行元素乘法來計算。通常情況下,我們可以通過特征圖中的每個像素向量計算出該像素標量的注意值。當涉及多個語義類別時,可以學習多維的注意力系數,使得每個注意力門能夠關注目標結構的1個子集,讓每個像素使用1個門控向量來確定聚焦區域并采用加性注意力方法來計算門控系數。通過這些方式,我們能夠建立適應不同場景和任務的注意力模式,并提高關鍵特征表達的準確性,加性注意力的公式如式(4)所示。

(4)

式中σ1表示激活函數,σ2表示Sigmoid函數,Θatt表示注意力門的一組參數,Wx、Wg和Ψ表示卷積操作,bΨ和bg表示對應卷積的偏置項。

3.2.2 將注意力門模塊應用于U-Net

在標準的U-Net架構中融合了前文提出的注意力門模塊,通過突出跳躍連接中的顯著特征,將粗尺度提取的信息應用于門控中,消除跳躍連接產生的不相關和嘈雜的響應。在前向傳播和反向傳播過程中,使用注意力門對神經元的激活進行過濾和減少產生背景區域梯度的權重等方法,確保更新底層模型參數時更專注于與給定任務相關的空間區域。具體實現的結構如圖6所示,對于l-1層的卷積參數,其更新方式如式(5)所示。

圖6 Attention U-Net網絡結構圖

(5)

4 試驗

4.1 試驗環境

本試驗采用Windows 10操作系統,11th Gen Intel(R) Core(TM) i5-11260H @ 2.60GHz (12 CPUs)處理器,NVIDIA GeForce RTX 3060顯卡,深度學習框架為Pytorch1.3.0版本,采用CamVid(Cambridge-driving Labeled Video Database)數據集。從中隨機挑選700張大小為970×720像素的圖像測試數據作為試驗的數據集,并按照5∶2的比例將數據集劃分為訓練集和測試集,訓練集共500張,測試集共200張,通過生成對抗網絡將數據集擴充到1 000張。

4.2 訓練參數及評價指標

4.2.1 訓練參數

使用本文的模型與U-Net+ +、FCN、DeepLabV1對擴充后的數據集進行訓練,將Batch size(1次訓練所抓取的數據樣本數量)設置為4,迭代次數設置為100,當損失趨于穩定時停止訓練。

4.2.2 評價指標

在深度學習領域,常用的語義分割評價指標有準確率(Accuracy)和損失值(Loss)等,這些參數可以較好地反映出一個模型的性能。準確率是用于評估分類模型的指標,簡單來講就是模型預測正確的數量占總量的比例,計算過程如式(6)所示。損失值是用于評價模型在訓練集上的擬合程度的指標,通常情況下,損失越小則訓練數據越擬合。

(6)

式中TP表示真正例,TN表示真負例,FP表示假正例,FN表示假負例。

4.3 試驗結果及分析

從圖7可以看出隨著時間的增長,在80輪以后,準確率和損失值逐漸趨于穩定,結合表1可以得出在第100輪時準確率約為95.4%。從表2各種模型訓練100輪的結果可以看出,在各自損失值趨于穩定的情況下,與其他模型相比,本文模型的訓練準確率較高,損失值與U-Net++、DeepLabV1較為接近。

圖7 訓練結果圖

表1 對模型進行100輪訓練的結果

表2 與其他模型的對比結果

5 結語

針對當前深度神經網絡在圖像分割過程中普遍存在的處理耗時長、分割準確率不高的問題,本文在使用生成對抗網絡擴充數據集后,提出了一種融入注意力機制的U-Net模型用于圖像分割。具體來講,就是以U-Net網絡為基本框架,在每個解碼器層中都添加1個注意力門模塊,幫助模型更準確地區分前景和背景,提高學習效率,加強特征區分。通過對比試驗得出,相較于各傳統語義分割模型,如U-Net++、SegNet等,本文的模型訓練時間較短,同時又具有較高的分割準確率。