基于特征累積的語義分割網(wǎng)絡(luò)設(shè)計(jì)

史正一，孫力

（華東師范大學(xué)通信與電子工程學(xué)院，上海 200241）

語義分割作為機(jī)器視覺的基本任務(wù)之一，在自動(dòng)駕駛、智能城市、醫(yī)療診斷等多個(gè)領(lǐng)域有著廣泛的應(yīng)用。近年來，隨著深度學(xué)習(xí)技術(shù)的發(fā)展，研究者們開始廣泛地使用卷積神經(jīng)網(wǎng)絡(luò)（Convolutional Neural Networks,CNNs）和視覺注意力編碼器（Vision Transformer）提取深度特征，進(jìn)而完成圖像的分割工作[1-2]。

現(xiàn)有的主流方法一般使用CNNs和VisionTransformer作為骨干網(wǎng)絡(luò)（Backbone），之后設(shè)計(jì)不同的連接器（Head），計(jì)算并強(qiáng)化圖像特征空間的依賴關(guān)系[3-8]。然而，值得注意的是：自然圖像中相同類別的特征往往存在類似的分布，即同一類別的物體，在不同圖像中的特征分布仍然是相似的。而現(xiàn)有的方法過度依賴于當(dāng)前圖像，并不能很好地結(jié)合數(shù)據(jù)集中其余圖像的信息。

針對(duì)以上討論，該文設(shè)計(jì)了一種基于動(dòng)量累積的語義分割網(wǎng)絡(luò)，該網(wǎng)絡(luò)能有效地捕獲并提取數(shù)據(jù)集的類別分布特征。與現(xiàn)有的主流方法相比，可以取得更高的分割精確度。

1 網(wǎng)絡(luò)設(shè)計(jì)

1.1 概述

該網(wǎng)絡(luò)的結(jié)構(gòu)如圖1 所示。利用全卷積連接器（Fully Convolutional Network Head,FCNHead）融合骨干網(wǎng)絡(luò)不同階段（Stage）的特征，得到初步分割結(jié)果的預(yù)測(cè)。根據(jù)該預(yù)測(cè)，使用類中心計(jì)算模塊統(tǒng)計(jì)特征空間中不同類別的數(shù)據(jù)分布，并在訓(xùn)練過程中使用動(dòng)量累積模塊逐步更新該分布值，之后將其存儲(chǔ)在類中心存儲(chǔ)單元中。測(cè)試階段，網(wǎng)絡(luò)通過余弦相似度的對(duì)比，選擇對(duì)應(yīng)的類中心，并使用注意力算法對(duì)圖像特征進(jìn)行更新，據(jù)此得到最終的預(yù)測(cè)結(jié)果。

圖1 網(wǎng)絡(luò)結(jié)構(gòu)示意圖

1.2 類中心計(jì)算模塊

考慮到自然圖像分割數(shù)據(jù)集中，不同樣本之間的空間結(jié)構(gòu)存在較大差異，而通道維度的類別特征信息則有一定的相似性。類似于ACFNet[9]、OCRNet[10]等Coarse-to-Fine 框架，該模塊使用受監(jiān)督的初步預(yù)測(cè)結(jié)果，統(tǒng)計(jì)不同通道的特征分布，即類中心。

對(duì)于k分類的分割任務(wù)，給定骨干網(wǎng)絡(luò)不同階段的輸出X={X1,X2,X3,X4}，可以采用FCNHead 完成特征融合，該過程可以表示為：

式中，ψ指上采樣函數(shù)，它可以將不同大小的輸入特征通過雙線性插值的方式，形變到相同的空間尺度。Proj 指映射函數(shù)，其使用1×1 卷積的方式，將不同輸入映射成具有相同的通道數(shù)的特征。[·]指Concate 函數(shù)，它會(huì)在通道的維度上對(duì)輸入的特征進(jìn)行堆疊。在此之后，MLP 指1×1 的卷積，它可以將輸入映射為k個(gè)通道的分類結(jié)果。

對(duì)于當(dāng)前圖像的類中心，可以使用如下方法進(jìn)行統(tǒng)計(jì)：

其中，∈?C×HW為X4的轉(zhuǎn)置，C代表骨干網(wǎng)絡(luò)提取到的X4的通道數(shù)。這里使用Softmax 函數(shù)δ在空間尺度的方向（即HW方向）進(jìn)行歸一，之后使用矩陣乘法的方式與下采樣后的Y＇進(jìn)行空間方向的加權(quán)融合，統(tǒng)計(jì)得到了該圖的類中心Center ∈?C×k。

1.3 單一子類的類中心存儲(chǔ)單元

為了讓網(wǎng)絡(luò)在預(yù)測(cè)階段盡可能地考慮到數(shù)據(jù)集其他圖片的特征分布，該方法引入了獨(dú)特的類中心存儲(chǔ)單元U。該單元不隨網(wǎng)絡(luò)的梯度變化而更新，相反地，在參數(shù)更新階段，使用動(dòng)量累積的方式，對(duì)每張圖每個(gè)類別的分布進(jìn)行滑動(dòng)平均，并將該平均值存儲(chǔ)在U中。在預(yù)測(cè)階段，存儲(chǔ)單元U將代表數(shù)據(jù)集不同類別特征分布的固定值，獨(dú)立地參與后續(xù)運(yùn)算。

該模塊的輸入為類中心計(jì)算模塊的輸出Center，輸出為維度相同的、累積得到的數(shù)據(jù)集類中心Z∈?C×k。初始化矩陣U∈?C×k，該矩陣不根據(jù)梯度隨機(jī)下降的方式進(jìn)行參數(shù)更新，其主要目的是記錄類中心逐步收斂后的均值，對(duì)于訓(xùn)練過程中的任意第t次迭代，此時(shí)有：

式中，μ為動(dòng)量系數(shù)，為大于0或小于1的超參數(shù)，用來平衡更新的速度。該值越大，U更新的幅度越小。

在推測(cè)階段，類中心存儲(chǔ)單元的值不再進(jìn)行改變。它將向網(wǎng)絡(luò)直接提供累積得到的平均值，在注意力模塊中完成對(duì)深層特征的權(quán)重分配。

1.4 多子類的類中心存儲(chǔ)單元

以上的類中心存儲(chǔ)和應(yīng)用方式在實(shí)際的推測(cè)過程中存在一定的弊端：由于自然圖像的復(fù)雜性，同一類物體的分布中可能存在若干個(gè)子類。以汽車為例，可能包含越野車、轎車等不同的子類別，其分布雖然相似但是仍有明顯的不同。此外，在特定的訓(xùn)練輪次中，不一定所有的類別都出現(xiàn)在圖像內(nèi)，因此，對(duì)于未出現(xiàn)的類別，其統(tǒng)計(jì)類中心可以視為噪聲。在訓(xùn)練和預(yù)測(cè)階段，如果對(duì)于每一個(gè)類別僅累積一個(gè)類中心，就可能會(huì)導(dǎo)致中心的失準(zhǔn)。為了解決以上問題，使用擴(kuò)充類中心的方法，進(jìn)一步完善網(wǎng)絡(luò)的結(jié)構(gòu)。

圖2 和圖3 以兩個(gè)子類中心為例，描述了多中心的情況下，存儲(chǔ)單元U在訓(xùn)練和測(cè)試階段的工作模式。

圖2 雙子類中心的訓(xùn)練模式

圖3 雙子類中心的測(cè)試模式

假設(shè)為每一個(gè)類別設(shè)計(jì)了p個(gè)子類中心，網(wǎng)絡(luò)需要初始化p個(gè)存儲(chǔ)單元，記為{U1,U2,…,Up}。如圖2 所示，在訓(xùn)練階段，對(duì)于輸入本圖類中心，需要計(jì)算其與每一個(gè)子類中心的余弦相似度。假設(shè)第q個(gè)子類中心的相似度為Sq，有：

得到的Sq∈?k，代表第q個(gè)子類中心在k個(gè)類別上的得分，故有：

據(jù)此可以得到每一個(gè)類別在不同子類上的得分，得分越高則本圖類中心與其對(duì)應(yīng)的子類中心的相似度越高，本圖中的物體更可能屬于該子類。因此，需要使用當(dāng)前圖像的統(tǒng)計(jì)結(jié)果，更新對(duì)應(yīng)的子類中心。

訓(xùn)練階段，根據(jù)S在p維度上的最大值索引，網(wǎng)絡(luò)可以有目的地更新對(duì)應(yīng)的子類中心，而其他的子類中心則維持不變。這種方法可以有效地避免因?yàn)槲闯霈F(xiàn)類別的原因而產(chǎn)生的噪聲，也可以適應(yīng)不同的子類，從而全面地描述數(shù)據(jù)集中的特征分布特點(diǎn)。

測(cè)試階段，同樣使用公式（4）和公式（5）的方法計(jì)算得分S，并根據(jù)最大值的索引從{U1,U2,…,Up}∈?k×C×p中選取最相似的中心，并拼接得到該模塊的輸出結(jié)果U∈?k×C。

1.5 基于類中心的注意力算法

近年來，注意力算法被廣泛地應(yīng)用到語義分割的相關(guān)任務(wù)中。該方法通過計(jì)算協(xié)方差矩陣的方法，衡量映射后的特征單元（Token）之間的相似度，據(jù)此重新分配其權(quán)重。

由于類中心存儲(chǔ)單元給出的數(shù)據(jù)集類中心描述了當(dāng)前數(shù)據(jù)集下，特征通道和類別的對(duì)應(yīng)關(guān)系。因此，它能夠有效地指導(dǎo)網(wǎng)絡(luò)進(jìn)一步地提取深層特征的信息。該文使用交叉注意力算法實(shí)現(xiàn)這一過程：

對(duì)于骨干網(wǎng)絡(luò)提取到的深層特征X4∈?HW×C和類中心存儲(chǔ)單元給出的數(shù)據(jù)集類中心U∈?k×C，使用映射函數(shù)將其調(diào)整為輸入鍵值Query(Q)、Key(K) 和Value(V)：

式中，WQ∈?C×C、WK∈?C×C和WQ∈?C×C為參數(shù)獨(dú)立的矩陣；對(duì)于深層特征X4，在不同的像素空間上共享權(quán)重；而對(duì)于類中心U，其不同類別共享權(quán)重。該處使用三組1×1 卷積來實(shí)現(xiàn)這一過程。

之后，使用矩陣乘法計(jì)算Query 中每個(gè)空間位置上對(duì)k個(gè)類別的響應(yīng)A：

式中A∈?HW×k。同時(shí)，V描述了類別和通道的對(duì)應(yīng)關(guān)系，因此可以使用如下方式，將A對(duì)齊到樣本空間：

δ同樣代表softmax 函數(shù)，在規(guī)范化權(quán)重的同時(shí)引入非線性表達(dá)，d為正則化因子，是固定的常數(shù)。同時(shí)，類似于Non-local Network[12]等工作，該模塊引入了shortcut 的機(jī)制，聯(lián)立公式（6）-（8），整個(gè)的注意力模塊可以描述為：

需要強(qiáng)調(diào)的是，網(wǎng)絡(luò)需要阻止節(jié)點(diǎn)U上的梯度回傳。訓(xùn)練階段，U只根據(jù)動(dòng)量累積更新，而測(cè)試階段，當(dāng)前圖像類中心僅用于選擇合適的數(shù)據(jù)集子類中心，自身則不直接參與注意力模塊的計(jì)算。

隨后，網(wǎng)絡(luò)根據(jù)特征F，利用FCN 連接器直接預(yù)測(cè)出最終的分割結(jié)果：

1.6 損失函數(shù)

由于該方法輸出了兩個(gè)分割圖，即初步分割結(jié)果Y＇和最終分割結(jié)果Y，網(wǎng)絡(luò)需要調(diào)整二者的損失函數(shù)權(quán)重來完成監(jiān)督。此處統(tǒng)一使用Softmax Cross Entropy(CE)函數(shù)來約束分割結(jié)果，并使用超參數(shù)λ進(jìn)行平衡，設(shè)圖像的標(biāo)注信息（Ground Truth）為GT，有：

2 實(shí) 驗(yàn)

2.1 環(huán)境配置

以下實(shí)驗(yàn)的代碼環(huán)境基于PyTorch 進(jìn)行搭建，使用4 張GTX2080 GPU 進(jìn)行訓(xùn)練和測(cè)試。為了驗(yàn)證模型的有效性，所有實(shí)驗(yàn)均在通用數(shù)據(jù)集Cityscapes 和Pascal Context 上進(jìn)行了測(cè)試。

在使用CNNs（ResNet）作為骨干網(wǎng)絡(luò)時(shí)，模型使用初始學(xué)習(xí)率為0.01的隨機(jī)梯度下降（SGD+Momentum）進(jìn)行梯度回傳和更新；在使用Vision Transformer（DeiT[13]）作為骨干網(wǎng)絡(luò)時(shí)，模型使用初始學(xué)習(xí)率為6×10-5的AdamW 優(yōu)化器進(jìn)行更新。模型的骨干網(wǎng)絡(luò)均在ImageNet-1k 數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練。

在模型訓(xùn)練的階段，為了避免過擬合，使用了如下的數(shù)據(jù)增強(qiáng)方法：首先，將圖像隨機(jī)裁剪至同樣大小的分辨率，在隨機(jī)水平翻轉(zhuǎn)后，加入隨機(jī)的高斯噪聲和飽和度進(jìn)行調(diào)整。對(duì)于Pascal Context 數(shù)據(jù)集，每次采樣16 個(gè)樣本作為一個(gè)批次，共進(jìn)行8 萬次迭代；對(duì)于Cityscapes 數(shù)據(jù)集，每次采樣8 個(gè)樣本作為一個(gè)批次，共進(jìn)行8 萬次迭代。

2.2 對(duì)比實(shí)驗(yàn)

該節(jié)在CNNs和VisionTransformer 下，對(duì)比所提方法與主流語義分割網(wǎng)絡(luò)的結(jié)果。使用平均準(zhǔn)確度（mAcc）和平均交并比（mIoU）作為指標(biāo)進(jìn)行性能評(píng)估。為保證對(duì)比公平，所有模型均在相同的條件下在訓(xùn)練集上完成訓(xùn)練。測(cè)試階段均使用驗(yàn)證集，并且，不使用包括多尺度和水平翻轉(zhuǎn)等數(shù)據(jù)增強(qiáng)效果來提升性能。

表1和表2對(duì)比了Pascal Context 下，以ResNet和DeiT 作為骨干網(wǎng)絡(luò)，不同連接器的性能指標(biāo)。該方法相對(duì)于主流算法，可以有效地提升分割性能；對(duì)比同樣采用通道維度注意力算法的OCRNet，在ResNet下mIoU 可以提升約0.34%，在DeiT 下mIoU 可以提升約0.10%，與PSPNet、DeepLab、UperNet[14-16]等強(qiáng)化空間感受野[17-18]的算法相比，性能上也有優(yōu)勢(shì)，比較SOTA方法，也有一定的提升；表3則在Cityscapes數(shù)據(jù)集下進(jìn)行了對(duì)比，可以看到該方法在不同的數(shù)據(jù)集上仍保持有效，以ResNet 作為骨干網(wǎng)絡(luò)，該方法相對(duì)于OCRNetmIoU 可以提升約0.34%。

表1 ResNet下Pascal Context的對(duì)比結(jié)果

表2 DeiT下Pascal Context的對(duì)比結(jié)果

表3 ResNet下Cityscapes的對(duì)比結(jié)果

3 結(jié)束語

該文設(shè)計(jì)了一種新型的語義分割網(wǎng)絡(luò)，在訓(xùn)練過程中可以動(dòng)態(tài)地累積每一個(gè)類別的數(shù)據(jù)分布。新建了類中心存儲(chǔ)單元，并在數(shù)據(jù)集的層面上得到所有類別的均值，結(jié)合實(shí)際情況，討論了不同的類中心設(shè)計(jì)方案。在網(wǎng)絡(luò)的預(yù)測(cè)階段，該方法利用當(dāng)前圖像的統(tǒng)計(jì)結(jié)果，在類中心存儲(chǔ)單元中選擇最接近的節(jié)點(diǎn)，并對(duì)深層特征進(jìn)行權(quán)重的重分配。實(shí)驗(yàn)證明，該方法能夠適應(yīng)不同的骨干網(wǎng)絡(luò)，在Pascal Context和Cityscapes 上獲得性能的提升。

然而，該方法在計(jì)算類中心時(shí)，過分強(qiáng)調(diào)了初步預(yù)測(cè)結(jié)果，對(duì)骨干網(wǎng)絡(luò)的性能有較高的要求，這導(dǎo)致整個(gè)模型的參數(shù)量和計(jì)算量較大，不利于實(shí)際的部署。在后續(xù)的研究中，將會(huì)進(jìn)一步探究模型在輕量化網(wǎng)絡(luò)中的潛能。