基于級聯注意力和邊界預測改進的輕量Segformer語義分割

關鍵詞：Ｓｅｇｆｏｒｍｅｒ；級聯注意力；輔助任務；邊界預測；梯度手術

中圖分類號：ＴＰ３９１． ４ 文獻標志碼：Ａ 開放科學（資源服務）標識碼（ＯＳＩＤ）：

文章編號：１００３－３１０６（２０２４）１１－２５８５－０９

０引言

圖像的語義分割任務是對圖像中的每個像素按照區域劃分，分為不同的類別，是計算機視覺的一項重要任務。目前，圖像語義分割許多下游任務中有很多應用，如醫學圖像分割、自動駕駛等。圖像的語義分割是當前計算視覺領域的一個重要研究方向。

在２０１５ 年，Ｌｏｎｇ 等［１］首次將深度學習應用到語義分割任務并提出了全卷積網絡（Ｆｕｌｌｙ Ｃｏｎｖｏｌｕ-ｔｉｏｎａｌ Ｎｅｔｗｏｒｋ，ＦＣＮ），實現了基于卷積神經網絡（Ｃｏｎｖｏｌｕｔｉｏｎ Ｎｅｕｒａｌ Ｎｅｔｗｏｒｋ，ＣＮＮ）模型的語義分割模型，之后出現了一系列如Ｄｅｅｐｌａｂ 系列［２］、ＰＳＰ-Ｎｅｔ［３］、ＳｅｇＮｅｔ［４］等對卷積網絡的改進語義分割方法。２０２０ 年，Ｇｏｏｇｌｅ 團隊提出了適用于計算機視覺的Ｖｉｓｉｏｎ Ｔｒａｎｓｆｏｒｍｅｒ 網絡，并在目標檢測和語義分割任務中取得了良好的效果，超越了當時大部分的ＣＮＮ 方法。后續出現了許多基于Ｔｒａｎｓｆｏｒｍｅｒ 架構的語義分割方法。Ｗａｎｇ 等［５］將金字塔結構加入到ＶｉＴ 網絡模型架構中，使網絡獲取到提取金字塔特征的能力。通過提取到金字塔特征，使網絡在與ＶｉＴ 網絡相同參數量下，能夠取得更高的語義分割效果。Ｌｉｕ 等［６］為解決ＶｉＴ 網絡計算量大的問題提出Ｓｗｉｎ Ｔｒａｎｓｆｏｒｍｅｒ 架構，用滑動窗口注意力替換其中的自注意力，在減少網絡參數量的同時，提高網絡在圖像上的預測效果。Ｄｏｎｇ 等［７］則在Ｓｗｉｎ Ｔｒａｎｓ-ｆｏｒｍｅｒ 的基礎上，將滑動窗口替換為十字窗口注意力，進一步優化了網絡的參數量和計算量。最后，為拓展Ｔｒａｎｓｆｏｒｍｅｒ 架構在移動端的進一步應用，Ｘｉｅ等［８］提出了基于Ｔｒａｎｓｆｏｒｍｅｒ 的Ｓｅｇｆｏｒｍｅｒ 網絡模型，通過對自注意力進行改進以及使用了混合全連接層減少原有Ｔｒａｎｓｆｏｒｍｅｒ 的參數量，從而減少了網絡的參數量和計算復雜度。從上述研究中可以看到，Ｔｒａｎｓｆｏｒｍｅｒ 網絡模型存在計算參數量大、復雜度高的問題。因此，減少Ｔｒａｎｓｆｏｒｍｅｒ 網絡的參數量并降低網絡的復雜度，成為當前圖像領域對Ｔｒａｎｓｆｏｒｍｅｒ 網絡的研究熱點問題之一。

鑒于Ｔｒａｎｓｆｏｒｍｅｒ 架構在移動端受到參數量和計算量的限制，由Ｍａ 等［９］的研究可以看出，任務的添加會出現梯度沖突問題。因此，本文基于Ｓｅｇｆｏｒｍｅｒ網絡設計了一種新的基于級聯注意力和邊界預測的輕量語義分割網絡，主要工作如下：

① 對于原有的輕量Ｓｅｇｆｏｒｍｅｒ解碼器進行了改進，加入了級聯注意力機制，使得網絡可以聚合并加強多尺度特征。

② 為解決原網絡對邊界信息提取能力弱的問題，通過加入簡單的輕量邊界預測輔助任務，并采用訓練階段啟動，測試階段關閉的策略，在減少網絡參數量的同時，提高網絡對邊界信息的感知能力。

③ 訓練階段，為防止因任務添加而導致出現梯度沖突，進而導致訓練效果下降的問題，在訓練部分加入了梯度手術優化算法，減小了任務之間間斷的梯度沖突問題，提高了模型的訓練效果。

１改進的輕量語義分割網絡

首先介紹網絡的整體架構，然后介紹改進的基于級聯注意力機制的語義分割解碼器，最后介紹設計的輕量邊界預測輔助任務和改進后的訓練策略。

１．１整體網絡架構

整體網絡架構如圖１ 所示，主要采用了編碼器和解碼器的結構。對于輸入的圖像，首先使用Ｓｅｇ-ｆｏｒｍｅｒ 的主干網絡Ｍｉｔｂ０ 對圖像進行特征提取工作。通過使用自注意力、全連接層和重疊補丁合并的操作，分別提取出相對于原圖１／４、１／８、１／１６、１ ／ ３２ 大小的Ｆ１、Ｆ２、Ｆ３、Ｆ４ 特征圖，之后使用１ ×１的卷積進行通道聚合，分別送入語義分支和邊界分支進行預測工作。

語義分支主要由改進的多尺度分割解碼器組成，對輸入的多尺度特征，使用門注意力、卷積注意力和上采樣組成的級聯注意力機制，通過從下到上的方式增強每層的特征信息，使得每層的特征都能更好地捕捉到深層特征中的關鍵信息。在邊界分支階段，對傳入的多尺度特征，使用卷積層和全連接層進行邊界預測，并計算損失函數。整體網絡的訓練策略為訓練時啟用邊界預測進行輔助訓練，在測試階段隱藏輔助任務的參數量，通過這種訓練策略在不增加任何多余參數量的基礎上，增強網絡對邊界信息的感知能力。

１．２改進的多尺度語義分割解碼器

原有的解碼器結果因過于簡單而無法有效利用多尺度信息，本文加入了級聯注意力［１０］用于改進原有的語義分割解碼器。

改進后的多尺度語義分割解碼器整體分為４ 層，每層提取到的特征分別為Ｆ１、Ｆ２、Ｆ３、Ｆ４，對應的尺度為原圖的１／４、１／ ８、１／１６、１／３２。對于提取的深層特征Ｆ４，首先使用一個１×１ 的卷積對通道進行聚合，然后使用卷積注意力模塊（Ｃｏｎｖｏｌｕｔｉｏｎａｌ Ａｔ-ｔｅｎｔｉｏｎ Ｍｏｄｕｌｅ，ＣＡＭ）提取空間和通道信息。將使用注意力機制加強后的特征送入上一層并在當前層中使用１×１ 的卷積進行通道變換。

在上一層中，對傳入的加強后深層特征，通過使用門注意力機制計算深層特征與當前層特征的相似度并映射到當前層特征中，使當前層特征Ｆ３ 聚合到深層特征信息。將聚合后的特征與傳入深層特征按照通道進行拼接工作，并繼續使用ＣＡＭ，從空間和通道角度對當前層特征進行特征加強。重復此過程，從下而上，使得深層特征信息聚合到了淺層特征中，豐富了網絡整體信息。將加強后的多尺度特征按照通道進行拼接，使用全連接層進行語義分割的預測。

級聯注意力中主要使用了門注意力模塊和ＣＡＭ。門注意力如圖２ 所示，借助門注意力層對傳過來的低分辨率特征與高分辨率特征結合計算出一張權重圖，然后將權重圖與當前層特征點乘，得到一張融合了深層特征的特征圖。

卷積注意力塊主要由通道注意力（Ｃｈａｎｎｅｌ Ａｔｔｅｎｔｉｏｎ，ＣＡ）、空間注意力（Ｓｐａｔｉａｌ Ａｔｔｅｎｔｉｏｎ，ＳＡ）和２層３×３卷積串聯組成。ＣＡＭ 首先通過ＣＡ 對輸入特征的通道部分，通過加強重要的通道，抑制不感興趣的通道進行特征加強。接著使用ＳＡ在空間區域上加強感興趣的特征，抑制不感興趣的特征。最后，使用了２ 層卷積進一步聚合前面的特征。

１．３多尺度輕量邊界輔助

本文與當前邊界任務對語義分割任務的處理做法不同，如圖３ 所示，大部分工作都將邊界預測到的信息加入到了語義分割任務中，而這樣會導致出現多余的邊界參數，增大網絡參數量。因此本文借助輔助任務思想，設計了一個輕量邊界輔助在訓練階段借助參數更新，使網絡獲得邊界感知能力，并優化了分割邊界。

本文參考了原有解碼器，將其用作邊界標簽預測，由于主干網絡本身進行過預訓練，所以添加的簡單邊界預測模塊可以快速提高預測的準確性，更好地推導網絡的參數更新。

其中語義分割標簽范圍為０～２５５，對應的邊界標簽設置為０ 和２５５，以加強２ 個任務之間的聯系。

邊界預測模塊整體結構如圖４ 所示。首先，通過使用４ 個１×１ 的卷積改變每一層特征通道數，然后進行上采樣操作，提高特征的分辨率。將特征沿著通道進行拼接，之后經過一個由１×１ 卷積和歸一層組成的多層感知機（Ｍｕｌｔｉ-Ｌａｙｅｒ Ｐｅｒｃｅｐｔｉｏｎ，ＭＬＰ）層進行邊界預測。整體結構與Ｓｅｇｆｏｒｍｅｒ 的解碼器類似，從實驗中可以看到，加入的邊界預測模塊由于輕量的結構，避免了因訓練初期繁瑣的網絡結構帶來的訓練損失負擔而導致的訓練性能下降。

１．４改進的訓練策略

由于邊界輔助任務的增加，會使網絡變為一個多任務網絡，隨之帶來了梯度沖突的問題。對此，本文設計了一個適用于添加邊界輔助任務的訓練策略，即在訓練階段加入梯度手術優化［１１］算法，用于減少加入邊界輔助帶來的梯度沖突問題。

本文使用交叉熵損失函數用于損失函數計算。交叉熵損失函數是多任務中常用的損失函數，通常伴隨ｓｏｆｔｍａｘ 函數使用，具體計算如式（６）所示。整體網絡損失函數分別由語義分割損失和邊界輔助損失組成。

式（８）表示將沖突的語義分割任務梯度向量投影到任務之間的法平面上。

④ 按照②、③更新每一個梯度向量

本文在梯度更新中加入梯度手術優化算法，雖增大了顯存占用（如表１所示），但卻優化了網絡訓練過程中因加入輔助任務而產生的梯度沖突干擾，使得網絡訓練結果更好。

使用單卡１６批次訓練，可以看到，添加梯度手術后增加的顯存相比原顯存占用并不多。本文網絡通過增加了１５５４ ＭＢ的顯存占用，使得整個網絡訓練效果更好。

２實驗

２．１數據集

ＡＤＥ２０Ｋ 是一個１５０ 語義類的場景數據集，分別由２０ ０００ 張訓練圖片和２ ０００ 張驗證圖片組成，圖片類型為ＲＧＢ 三通道圖像。Ｃｉｔｙｓｃａｐｅｓ 是關于城市街道的語義數據集，擁有５ ０００ 張精細標注的圖像和２０ ０００ 個粗略標注的附加圖像。ＡＤＥ２０ｋ 數據集訓練時采用５１２ ｐｉｘｅｌ × ５１２ ｐｉｘｅｌ 的尺度，Ｃｉｔｙｓｃａｐｅｓ 數據集采用１ ０２４ ｐｉｘｅｌ×１ ０２４ ｐｉｘｅｌ 的尺度。其中邊界標簽的獲取可以參照文獻［１２］中的獲取方式。

２．２評估指標

本文采用平均像素精度（ａｖｅｒａｇｅ ＰｉｘｅｌＡｃｃｕｒａｃｙ，ａＡｃｃ）、類別平均像素準確率（ｍｅａｎ ＰｉｘｅｌＡｃｃｕｒａｃｙ，ｍＡｃｃ）和平均交并比（ｍｅａｎ Ｉｎｔｅｒｓｅｃｔｉｏｎｏｖｅｒ Ｕｎｉｏｎ，ｍＩｏＵ）來衡量，

２．３實驗環境

本文提出的網絡是在ｍｍｓｅｇｍｅｎｔａｔｉｏｎ 框架下實現的，使用的顯卡為ＮＶＩＤＩＡ ３０９０，主干網絡Ｍｉｔ-ｂ０是在Ｉｍａｇｅｎｅｔ-１ｋ 上預訓練的權重。

與Ｓｅｇｆｏｒｍｅｒ 訓練設置相同，學習率為６ ×１０^－５ ，權重衰減為０．０１，訓練批次為１６，在ＡＤＥ２０ｋ 上進行了１６０ ０００ 次迭代訓練。

２．４消融實驗

本文在ＡＤＥ２０ｋ 數據集上進行了消融實驗，表中Ｓｅｇｆｏｒｍｅｒ＋ＤＡ 表示只添加了級聯注意力模塊的模型，Ｓｅｇｆｏｒｍｅｒ＋ＢＰ 表示只添加了邊界預測輔助任務的模型，Ｓｅｇｆｏｒｍｅｒ＋ＤＡ＋ＢＰ ＋ＧＳ 是在基礎上添加級聯注意力和邊界輔助任務，在訓練階段添加了梯度手術的網絡模型。消融實驗如表２ 所示。改進的Ｓｅｇｆｏｒｍｅｒ 與原Ｓｅｇｆｏｒｍｅｒ 在ＡＤＥ２０ｋ 數據集可視化對比如圖６ 所示。

從表２ 可以看出，加入級聯注意力可以極大地提高Ｓｅｇｆｏｒｍｅｒ 網絡的性能，ｍＩｏＵ 和準確率都有所提高。添加級聯注意力模塊后，網絡可以更好地提取多尺度特征的信息。與原網絡相比，Ｓｅｇｆｏｒｍｅｒ＋BＰ的各項指標都有所提升，但提升并不明顯。

Ｓｅｇｆｏｒｍｅｒ＋ＤＡ ＋ＢＰ 是上述２個模塊結合的結果，效果明顯改善。與原網絡相比，ｍＩｏＵ增加了１．９３％ ，ａＡｃｃ增加了１．０３％ ，ｍＡｃｃ 增加了２． ３２％ 。表明級聯注意力增強了網絡提取多尺度特征的能力，邊界輔助任務使網絡能夠關注圖像的邊界特征。

添加邊界輔助任務后，任務之間可能會出現一些沖突，阻礙原有邊界預測的輔助效果。于是本文添加了梯度手術，解決了語義分割之間的梯度沖突，使得網絡訓練取得了更好的效果，其中ｍＩｏＵ增加了０．4５％ ，ａＡｃｃ 增加了０． ３６％ ，ｍＡｃｃ增加了０．５７％ 。

圖６中第１～ ４ 列分別是原始圖像、Ｓｅｇｆｏｒｍｅｒ 預測結果、Ｓｅｇｆｏｒｍｅｒ＋ＣＡ 預測結果和Ｓｅｇｆｏｒｍｅｒ＋ＣＡ ＋ＢＰ＋ＧＳ 預測結果。以第二行圖片為例，可以看出原來的Ｓｅｇｆｏｒｍｅｒ 網絡在預測結果中將海洋預測為天空，進行了錯誤預測。而在添加級聯注意力機制后成功預測出了結果。可以發現原有網絡在多尺度信息的聚合方面缺失，存在一定的缺陷。由第三行圖片可以看到，相比單獨添加級聯注意力，邊界輔助加入使得網絡預測到的邊界更為平整光滑，并且檢測到了一定的小目標信息，證明了添加邊界輔助的有效性。

２．５對比實驗

比較了添加損失梯度手術前后損失函數的變化，然后將本文的邊界輔助任務與其他研究中的邊界預測輔助任務進行比較。其中，ＡＢＰ［１３］邊界輔助表示了添加輔助注意力機制的邊界輔助任務，ＢＰ 代表本文的邊界輔助。在ＡＤＥ２０ｋ 測試集和Ｃｉｔｙｓｃａｐｅｓ 測試集上對本文的方法和其他語義分割方法進行了比較。

２．５．１增加了梯度手術前后的對比以及添加不同輔助任務的對比

添加梯度手術的前后Ｌｏｓｓ 變化和不同輔助任務對比如圖７所示。

由圖７（ａ）可以看出，添加了級聯注意力和邊界輔助的模型的損失函數曲線與原始Ｓｅｇｆｏｒｍｅｒ 網絡的損失函數曲線基本重合，但加入梯度手術操作后，網絡損失函數下降速度明顯比原本網絡更快。由此可以證明本文的訓練策略確實有效，提高了網絡的訓練效果，

為了驗證本文提出的輕量邊界輔助的有效性，對復雜邊界輔助任務與本文邊界輔助通過ｍＩｏＵ 和訓練階段的準確度進行對比。其中，Ｅｖｏｌｖｅ ＿Ｓｅｇｆｏｒｍｅｒ 表示Ｓｅｇｆｏｒｍｅｒ＋ＣＡ＋ＧＳ，指在添加了級聯注意力和梯度手術的情況，ＡＢＰ為復雜邊界輔助任務，其在邊界預測任務中添加了注意力機制，ＢＰ 表示本文輕量邊界預測任務，圖７（ｂ）則是對２ 種輔助任務在訓練時預測精度的對比曲線圖，每更新５０ 次參數就記錄一次總精確度。可以看出本文添加的邊界輔助收斂速度更快，能更快地對語義分割任務進行輔助。

不同輔助任務的對比如表３ 所示，可以看出，本文添加的輕量邊界輔助可以對語義分割任務進行輔助訓練。

２．５．２不同網絡的比較

表４ 和表５ 中對提出的模型在ＡＤＥ２０ｋ 測試集和Ｃｉｔｙｓｃａｐｅｓ 測試集上從參數大小、計算量和ｍＩｏＵ方面與其他模型進行了比較。

本文基于原有Ｓｅｇｆｏｒｍｅｒ 網絡設計了一種新的輕量網絡結構，繼承了原始網絡輕量的優點，并且具有更高的分割精度。從表４ 可以看出，雖然增加了２．６９ Ｍ 的參數量和２４． ６７ Ｇ 計算量，但相對于其他網絡，仍然具有輕量級的優勢。Ｌａｗｉｎ 是目前對Ｓｅｇｆｏｒｍｅｒ 解碼器最先進的改進。在ＡＤＥ２０ｋ 數據集上，本文方法在ｍＩｏＵ 上比Ｌａｗｉｎ 高１．３３％ ，在城市景觀上比Ｌａｗｉｎ 高１． ３５％ 。本文方法相比當前最先進的Ｓｅｇｆｏｒｍｅｒ 解碼器改進具有更高的語義分割精度，并且本文的方法在ＡＤＥ２０ｋ 數據集和Ｃｉｔｙｓｃａｐｅｓ數據集上都取得了很好的效果，證明了本文方法具有較好的魯棒性。

不同網絡可視化對比如圖８ 所示，可以看出，原方法存在著識別精度不高、邊界模糊的問題。從圖８ 第４ 行可以看出，餐盤被錯誤分割，并且邊界被識別模糊，而其他相同參數量的方法如ＰＳＰＮｅｔ 和ＤｅｅｐｌａｂＶ３＋都存在這些問題。第５ 列則是使用了ＯＣＲＮｅｔ 的方法，可以看到相比其他方法效果更好，但相對本文方法也有一些缺失，并且由表４ 可以看出ＯＣＲＮｅｔ 網絡比本文方法多出了３５７． ２２ Ｇ 的計算量和５．８２ Ｍ 的參數量。由此可以看到本文方法在輕量方面的優越性，更適用于移動端的進一步應用。

３結束語

本文借助注意力機制和邊界輔助設計了一種新的輕量語義分割架構，通過結合注意力機制和多尺度框架構建了多尺度級聯注意力解碼器。加入了設計的邊界輔助模塊，從梯度更新角度提高網絡的邊界感知能力并設計了一套適用本網絡的訓練策略，避免梯度沖突的同時，提高了網絡的訓練效率。對提出的改進網絡在ＡＤＥ２０ｋ 數據集上進行了消融實驗，驗證了各模塊的有效性；在不同數據集上進行的對比實驗彰顯了本文網絡的優勢。本文提出的基于級聯注意力和邊界預測的輕量改進Ｓｅｇｆｏｒｍｅｒ 網絡為其他語義分割網絡分割精度的提高和網絡的輕量化提供了一種新的思路，即通過相似任務，從梯度更新角度，提高網絡對其余特征的感知能力。改進的方法在輕量化方面還有一定的提升空間，對于注意力機制的添加而導致的參數量提升，后續會對其進一步優化，減少參數量的占用。此外，后期會考慮網絡在其他領域的進一步應用。

作者簡介

高延海 男，（１９９９—），碩士研究生。主要研究方向：語義分割、深度學習。

劉永帥 男，（１９７３—），中級實驗師。主要研究方向：知識圖譜。