Dynamic Unet+: 一種輕量精確的語義分割算法及應(yīng)用

陳朗 任洋甫 楊培 劉曉靜 崔亞超

關(guān)鍵詞：深度學(xué)習(xí)；卷積神經(jīng)網(wǎng)絡(luò)；圖像語義分割；語義分割算法應(yīng)用；自動(dòng)人像摳圖

0 引言

圖像語義分割是計(jì)算機(jī)對場景理解的重要環(huán)節(jié)，由于拍攝圖片通常會(huì)受到光照、角度、尺度以及拍攝圖片分辨率等多種復(fù)雜條件的影響，使圖像語義分割成了計(jì)算機(jī)視覺領(lǐng)域具有挑戰(zhàn)性的問題。例如像素聚類、閾值分割等傳統(tǒng)算法存在難以建立語義層級(jí)理解的問題，隨著2012年AlexNet在圖像分類領(lǐng)域取得的巨大成功，以卷積神經(jīng)網(wǎng)絡(luò)為代表的深度學(xué)習(xí)技術(shù)快速發(fā)展，同時(shí)也深刻影響了包括圖像語義分割在內(nèi)的眾多計(jì)算機(jī)視覺領(lǐng)域的研究。

隨著深度學(xué)習(xí)的引入，模型參數(shù)量也成倍增加，巨大的內(nèi)存消耗和計(jì)算開銷可能會(huì)成為算法應(yīng)用部署中的一大阻礙。大多數(shù)圖像語義分割算法具有非常龐大的參數(shù)量，例如在實(shí)驗(yàn)測試中，基于ResNet101實(shí)現(xiàn)的PSPNet[1]有高達(dá)70M的參數(shù)量，以VGG16為骨干網(wǎng)絡(luò)的FCN8s[2] 參數(shù)量更是超過了130M，以ResNet101 為骨干網(wǎng)絡(luò)實(shí)現(xiàn)的DeepLabv3+也有接近60M的參數(shù)量。這些網(wǎng)絡(luò)在擁有巨大參數(shù)量的同時(shí)，訓(xùn)練獲得原始論文中所報(bào)告的精度往往需要大量額外的訓(xùn)練數(shù)據(jù)以及工程訓(xùn)練技巧，缺乏一定的普適性和實(shí)用性。受O Ronneberger等人在醫(yī)學(xué)影像分割領(lǐng)域工作的啟發(fā)，本文利用輕量的Unet[3]的U形編解碼結(jié)構(gòu)和橫向越級(jí)連接，同時(shí)針對Unet對實(shí)際場景圖像中分割能力表現(xiàn)不足的問題進(jìn)行多方面的改進(jìn)，在對多個(gè)基準(zhǔn)數(shù)據(jù)集的測試中取得了很好的效果，且需要相對較少的訓(xùn)練技巧和數(shù)據(jù)量便可獲得較好的精度和結(jié)果。

本文的主要貢獻(xiàn)在以下兩個(gè)方面：

1） 提出了相對輕量、通用的動(dòng)態(tài)Unet 架構(gòu)，將ResNet融入U(xiǎn)net中組成動(dòng)態(tài)Unet，提高了模型對復(fù)雜場景特征的編碼能力，進(jìn)而提高了模型效果；

2） 替換了反卷積上采樣技術(shù)，并將像素重組[3]技術(shù)應(yīng)用到模型Dynamic Unet+中，進(jìn)一步提升了模型解碼器的能力并使模型獲得了一定的效果增益。

1 Unet 工作回顧和分析

1.1U 形編解碼器結(jié)構(gòu)

UNet（結(jié)構(gòu)如圖1 所示）最早發(fā)表在2015 年的MICCAI（醫(yī)學(xué)影像學(xué)術(shù)頂級(jí)會(huì)議）上，相比于FCN、PSPNet、DeepLab等發(fā)表在計(jì)算機(jī)視覺領(lǐng)域期刊和會(huì)議的工作來說，早期對其重視程度不高。UNet僅有7M左右的參數(shù)量，非常適用于應(yīng)用研究，同時(shí)作為一個(gè)靈活的語義分割架構(gòu)，具有很好的越級(jí)連接和編解碼器的方法。

在Unet基礎(chǔ)上也出現(xiàn)了許多UNet的改進(jìn)版本，例如UNet++[6]和UNet 3+[7]和基于Transformer的Tran?sUNet[8]等，但這些工作主要針對醫(yī)學(xué)影像分割，少有改進(jìn)工作是針對實(shí)際應(yīng)用場景的復(fù)雜圖片，且醫(yī)學(xué)影像和現(xiàn)實(shí)場景圖片在數(shù)據(jù)分布和復(fù)雜度方面存在很大差異，直接使用原始Unet進(jìn)行語義分割會(huì)存在域偏移等問題，導(dǎo)致模型效果較差，這也可能是早期Unet被圈內(nèi)學(xué)者低估的原因之一。

1.2 優(yōu)缺點(diǎn)分析

Unet相比其他語義分割模型，僅擁有7M左右的參數(shù)量，是作為研究的一個(gè)很好的基準(zhǔn)（Baseline），其中包含的越級(jí)連接和編解碼器思想使得模型蘊(yùn)藏著很大的潛力。然而由于醫(yī)學(xué)影像和自然場景圖片存在較大的數(shù)據(jù)分布變化，直接應(yīng)用會(huì)存在域偏移從而導(dǎo)致模型效果差，同時(shí)模型結(jié)構(gòu)不足以應(yīng)對復(fù)雜多變的自然圖像分割。

為了解決原始UNet在復(fù)雜自然圖像分割中能力欠缺的問題，如圖2所示，通過記錄ResNet等CNN系列網(wǎng)絡(luò)的下采樣模塊的位置更換了原始Unet的特征提取器，然后與解碼器對應(yīng)位置對齊，進(jìn)而復(fù)用了原始Unet的U形編解碼器結(jié)構(gòu)和橫向越級(jí)連接，得益于原始Unet的結(jié)構(gòu)優(yōu)點(diǎn)和ResNet在ImageNet的預(yù)訓(xùn)練權(quán)重，充分利用了遷移學(xué)習(xí)和微調(diào)的思想，使模型在此基礎(chǔ)上獲得了很好的效果增益。同時(shí)為了避免反卷積上采樣模塊可能存在的計(jì)算資源浪費(fèi)和性能次優(yōu)等問題，采用像素重組[3]進(jìn)行替換，增強(qiáng)了模型的解碼器能力，使得模型在恢復(fù)分割掩模時(shí)產(chǎn)生了更好的效果。

2 本文方法

圖2中展示了Dynamic Unet+動(dòng)態(tài)編碼器語義分割架構(gòu)。模型輸入是一個(gè)RGB的待分割圖片，模型輸出是一個(gè)與輸入同尺寸的類別注釋掩模圖，其中C為類別（Class）數(shù)量。其中沿用了Unet[5]的U形編解碼器結(jié)構(gòu)和橫向連接的基本結(jié)構(gòu)。

2.1 動(dòng)態(tài)編碼器

為了增強(qiáng)模型對于復(fù)雜多變的自然場景圖片編解碼器的能力和充分利用遷移學(xué)習(xí)技術(shù)（TransferLearning），如圖2所示則對原始的Unet的編碼器進(jìn)行了替換。在本文實(shí)驗(yàn)中使用ResNet系列網(wǎng)絡(luò)進(jìn)行代替，為了保持U形編解碼器結(jié)構(gòu)，將ResNet融入U(xiǎn)net的方式：去掉原始Unet的左半部的編碼器部分，記錄ResNet每一次下采樣的位置然后與右側(cè)的編碼器對應(yīng)位置進(jìn)行對齊。值得注意的是，對原始ResNet的最后的池化層及其池化之后的所有模塊進(jìn)行了截?cái)?，因?yàn)椴恍枰詈蟮膸讉€(gè)用來做分類的模塊，同時(shí)為了進(jìn)一步增強(qiáng)模型恢復(fù)分割細(xì)節(jié)的能力，將原始圖像與解碼器對應(yīng)分辨率的特征圖進(jìn)行了連接與合并（Connectand Concat），相比于原始的Unet，還在輸入和輸出部分直接進(jìn)行了越級(jí)連接，進(jìn)一步增強(qiáng)恢復(fù)分割細(xì)節(jié)的能力，為了進(jìn)一步增強(qiáng)特征解碼器的表現(xiàn)能力，使用更高的通道維度，例如1024、768等，并且將所有的上采樣模塊替換成了像素重組（Pixel Shuffle）模塊[3]。

2.2 像素重組上采樣技術(shù)

鄰域常用雙線性插值和反卷積進(jìn)行上采樣，由于雙線性插值通過直接計(jì)算得到上采樣圖，該過程并沒有可學(xué)習(xí)性，無法使用神經(jīng)網(wǎng)絡(luò)的強(qiáng)大學(xué)習(xí)能力進(jìn)行效果加持，同時(shí)反卷積上采樣模塊雖然有可學(xué)習(xí)的參數(shù)，在一定程度上增強(qiáng)了模型在上采樣過程的可學(xué)習(xí)性，促進(jìn)了模型的特征解碼和分割結(jié)果恢復(fù)，然而過多的零值填充使得卷積過程存在大量的計(jì)算浪費(fèi)。Wenzhe Shi等人[3]首次提出像素重組技術(shù)并成功應(yīng)用于圖像超分辨率生成，取得了很好的效果，本文創(chuàng)新性地將其應(yīng)用到了圖像語義分割，在幫助模型恢復(fù)分割細(xì)節(jié)和解碼特征方面，大大提高計(jì)算利用效率的同時(shí)，該技術(shù)表現(xiàn)出更好的特征解碼能力和分割效果，圖3展示了典型的反卷積和像素重組技術(shù)的計(jì)算過程示意圖。

2.3 損失函數(shù)

損失函數(shù)度量了模型的預(yù)測和真值標(biāo)簽之間的差距。實(shí)驗(yàn)中將圖像分割看作是像素級(jí)別的分類問題，因此本文使用的損失函數(shù)基于的多分類交叉熵?fù)p失函數(shù)：

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

3.1 實(shí)驗(yàn)數(shù)據(jù)集及評價(jià)指標(biāo)

為了驗(yàn)證本文方法的有效性，通過在PASCALVOC 2012和Camvid數(shù)據(jù)集上進(jìn)行了測試，報(bào)告的測試均是嚴(yán)格基于對應(yīng)的驗(yàn)證集并采用原始的訓(xùn)練集進(jìn)行訓(xùn)練。在PASCAL VOC2012數(shù)據(jù)集中，除了官方標(biāo)注，還報(bào)告了由第三方提供的含有1萬余張標(biāo)注信息的增強(qiáng)版本數(shù)據(jù)集上的測試結(jié)果，數(shù)據(jù)集詳情如表1所示。

評價(jià)指標(biāo)采用平均交并比（MIOU），即對每一類的預(yù)測和真值之間的交并比，然后在所有類別上求和取平均，設(shè)共有c + 1個(gè)類別（包括空類或背景類），則MIOU計(jì)算方式如下公式所示：

3.2 實(shí)驗(yàn)設(shè)置

實(shí)驗(yàn)中采用單張GTX1080Ti顯卡進(jìn)行訓(xùn)練，學(xué)習(xí)率均采用One Cycle策略，使用Adma優(yōu)化器，并采用大小1e-2的權(quán)重衰減，所使用的ResNet骨干特征提取網(wǎng)絡(luò)均采用其在ImageNet上的預(yù)訓(xùn)練權(quán)重以充分利用遷移學(xué)習(xí)技術(shù)，并使用多尺度訓(xùn)練的技巧：在PASCAL VOC 2012和Camvid兩個(gè)數(shù)據(jù)上，均先使用224×224的圖片進(jìn)行訓(xùn)練，收斂后在此基礎(chǔ)上將圖片輸入增加到448×448再進(jìn)行訓(xùn)練，然后再增加到512×512尺寸大小。受限于顯存等原因沒有繼續(xù)增加輸入圖像尺寸，報(bào)告的結(jié)果均是最后基于512×512進(jìn)行訓(xùn)練的最終結(jié)果。

3.3 實(shí)驗(yàn)結(jié)果

使用包含1萬余張的PASCAL VOC2012數(shù)據(jù)集訓(xùn)練了Dynamic Unet+ ResNet34模型，采用的驗(yàn)證集與原始的驗(yàn)證集保持一致，與其他模型的對比結(jié)果如表2所示：

值得注意的是，受限于實(shí)驗(yàn)環(huán)境和其他因素的影響，僅僅采用了最高ResNet34作為特征編碼器，在面對PASACL VOC2012這種相對復(fù)雜的分割任務(wù)來說可能是使得性能受限的重要原因之一，但仍然達(dá)到了81.8%的MIOU。這對于僅僅以ResNet34作為骨干網(wǎng)絡(luò)的模型來說已經(jīng)相當(dāng)可觀。

而且Dynamic UNet+在Camvid數(shù)據(jù)集上取得了非常具有競爭力結(jié)果。模型表現(xiàn)以及與SOTA模型的對比如表3 所示。為了平衡模型大小和精度，基于ResNet34骨干網(wǎng)絡(luò)進(jìn)行了實(shí)現(xiàn)，筆者的模型在Camvid數(shù)據(jù)集上取得了非常好的效果，在參數(shù)量保持較低的水平的情況下仍然保持了較高的分割精度。

3.4 消融實(shí)驗(yàn)

本節(jié)報(bào)告了Dynamic UNet+模型在PASCAL VOC2012數(shù)據(jù)集上相對于基準(zhǔn)模型UNet的精度變化和對比情況，同時(shí)探究了不同的編碼器嵌入和編碼器是否預(yù)訓(xùn)練對模型效果的影響，結(jié)果如表4所示。

可以看到，相對于基準(zhǔn)模型在PASCAL VOC 2012數(shù)據(jù)集上只取得了16%的MIOU值，本文的DynamicUNet+在將ResNet系列網(wǎng)絡(luò)嵌入后，模型效果大幅提升，顯著增強(qiáng)了模型處理復(fù)雜場景的能力，且呈現(xiàn)出隨著嵌入的編碼器網(wǎng)絡(luò)深度的增加，模型精度呈上升趨勢，同時(shí)在使用相同的特征編碼器的情況下，使用ImageNet上的預(yù)訓(xùn)練權(quán)重會(huì)顯著提升Dynamic UNet+的模型效果，這是本文模型能保持相對較低參數(shù)量的情況下仍能獲得如此具有競爭力的重要原因之一。

4 實(shí)驗(yàn)人像分割應(yīng)用

4.1 實(shí)驗(yàn)設(shè)置

為了進(jìn)一步驗(yàn)證模型的可行性，筆者測試了模型在人像分割方面的應(yīng)用能力。實(shí)驗(yàn)采用AutomaticPortrait Matting[16]和Aisegment 人像分割數(shù)據(jù)集，分別包含2 000張圖片和34 427張統(tǒng)一分辨率的高質(zhì)量的半身人像圖片。為了進(jìn)一步減少參數(shù)量，實(shí)驗(yàn)中選用了ResNet18作為特征提取的骨干網(wǎng)絡(luò)，在保證人像分割效果的同時(shí)，進(jìn)一步降低了模型參數(shù)量，模型配置及訓(xùn)練信息如表5所示：

根據(jù)Smith L等人的工作，采用One Cycle學(xué)習(xí)率策略。同時(shí)采樣與本文3.2節(jié)所述相同的參數(shù)設(shè)置和多尺度訓(xùn)練方法，使用Adma優(yōu)化器訓(xùn)練僅15個(gè)Ep?och模型基本收斂。

4.2 模型效果

為了測試模型的泛化能力，采用更貼近實(shí)際應(yīng)用場景的圖片進(jìn)行測試，筆者從互聯(lián)網(wǎng)上隨機(jī)選取了模型從未見過的肖像圖片進(jìn)行分割，分割結(jié)果如圖4所示，其中最左邊是模型輸入的原始RGB圖片，中間是模型生成的掩模，最右邊一列是對掩模圖進(jìn)行后處理制作的背景編輯應(yīng)用Demo，用于模擬真實(shí)的自動(dòng)證件照制作場景。

可以看到，基于ResNet18實(shí)現(xiàn)的Dynamic Unet+在處理簡單背景的肖像圖幾乎達(dá)到超越人類水平的結(jié)果，在較為復(fù)雜的背景下也依然保持較高水平。

5 結(jié)束語

UNet作為一個(gè)輕量靈活的醫(yī)學(xué)影像分割架構(gòu)早期并沒有受到機(jī)器學(xué)習(xí)圈內(nèi)學(xué)者的過多關(guān)注，并且少有工作將Unet遷移到復(fù)雜場景圖像語義分割。在本文中筆者充分利用的Unet的U形編解碼結(jié)構(gòu)和橫向越級(jí)連接，同時(shí)針對Unet在復(fù)雜語義和環(huán)境場景下表現(xiàn)力欠缺等問題對其進(jìn)行了多方面的改進(jìn)，提出了Dynamic Unet+，并舉例如何將ResNet嵌入到UNet模型中，該模型非常靈活，幾乎支持多種圖像分類或其他網(wǎng)絡(luò)的嵌入，例如使用Efficient Net將進(jìn)一步減少模型的參數(shù)量，提高速度。使用這些帶有預(yù)訓(xùn)練的模型作為骨干網(wǎng)絡(luò)，可以充分利用例如遷移學(xué)習(xí)、微調(diào)（Fine Tuning）等技術(shù)并從中受益。筆者的DynamicUnet+模型在PSCAL VOC2012、Camvid等數(shù)據(jù)集上取得了不錯(cuò)的效果，并且擁有相對更低的參數(shù)量和更快速度，促進(jìn)了語義分割在實(shí)際應(yīng)用上的落地。同時(shí)，通過對模型在人像分割應(yīng)用方面的初步測試，以較小的參數(shù)量和相當(dāng)?shù)木热〉昧撕芎玫膶?shí)際應(yīng)用效果，結(jié)果充分驗(yàn)證了本文的Dynamic Unet+算法的可行性和實(shí)用性。