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

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

關鍵詞：深度學習；卷積神經網絡；圖像語義分割；語義分割算法應用；自動人像摳圖

0 引言

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

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

本文的主要貢獻在以下兩個方面：

1） 提出了相對輕量、通用的動態Unet 架構，將ResNet融入Unet中組成動態Unet，提高了模型對復雜場景特征的編碼能力，進而提高了模型效果；

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

1 Unet 工作回顧和分析

1.1U 形編解碼器結構

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

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

1.2 優缺點分析

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

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

2 本文方法

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

2.1 動態編碼器

為了增強模型對于復雜多變的自然場景圖片編解碼器的能力和充分利用遷移學習技術（TransferLearning），如圖2所示則對原始的Unet的編碼器進行了替換。在本文實驗中使用ResNet系列網絡進行代替，為了保持U形編解碼器結構，將ResNet融入Unet的方式：去掉原始Unet的左半部的編碼器部分，記錄ResNet每一次下采樣的位置然后與右側的編碼器對應位置進行對齊。值得注意的是，對原始ResNet的最后的池化層及其池化之后的所有模塊進行了截斷，因為不需要最后的幾個用來做分類的模塊，同時為了進一步增強模型恢復分割細節的能力，將原始圖像與解碼器對應分辨率的特征圖進行了連接與合并（Connectand Concat），相比于原始的Unet，還在輸入和輸出部分直接進行了越級連接，進一步增強恢復分割細節的能力，為了進一步增強特征解碼器的表現能力，使用更高的通道維度，例如1024、768等，并且將所有的上采樣模塊替換成了像素重組（Pixel Shuffle）模塊[3]。

2.2 像素重組上采樣技術

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

2.3 損失函數

損失函數度量了模型的預測和真值標簽之間的差距。實驗中將圖像分割看作是像素級別的分類問題，因此本文使用的損失函數基于的多分類交叉熵損失函數：

3 實驗

3.1 實驗數據集及評價指標

為了驗證本文方法的有效性，通過在PASCALVOC 2012和Camvid數據集上進行了測試，報告的測試均是嚴格基于對應的驗證集并采用原始的訓練集進行訓練。在PASCAL VOC2012數據集中，除了官方標注，還報告了由第三方提供的含有1萬余張標注信息的增強版本數據集上的測試結果，數據集詳情如表1所示。

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

3.2 實驗設置

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

3.3 實驗結果

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

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

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

3.4 消融實驗

本節報告了Dynamic UNet+模型在PASCAL VOC2012數據集上相對于基準模型UNet的精度變化和對比情況，同時探究了不同的編碼器嵌入和編碼器是否預訓練對模型效果的影響，結果如表4所示。

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

4 實驗人像分割應用

4.1 實驗設置

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

根據Smith L等人的工作，采用One Cycle學習率策略。同時采樣與本文3.2節所述相同的參數設置和多尺度訓練方法，使用Adma優化器訓練僅15個Ep?och模型基本收斂。

4.2 模型效果

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

可以看到，基于ResNet18實現的Dynamic Unet+在處理簡單背景的肖像圖幾乎達到超越人類水平的結果，在較為復雜的背景下也依然保持較高水平。

5 結束語

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