基于一致性的半監(jiān)督語義分割方法

在本工作中，來自阿德萊德大學(xué)、烏魯姆大學(xué)的研究者針對當(dāng)前一致性學(xué)習(xí)出現(xiàn)的三個問題做了針對性的處理， 使得經(jīng)典的 teacher-student 架構(gòu) （A.K.A Mean-Teacher） 在半監(jiān)督圖像切割任務(wù)上得到了顯著的提升。

該研究已被計算機(jī)視覺頂會 CVPR 2022 大會接收，論文標(biāo)題為《Perturbed and Strict Mean Teachers for Semi-supervised Semantic Segmentation》：

背景

語義分割是一項重要的像素級別分類任務(wù)。但是由于其非常依賴于數(shù)據(jù)的特性（data hungary）， 模型的整體性能會因為數(shù)據(jù)集的大小而產(chǎn)生大幅度變化。同時， 相比于圖像級別的標(biāo)注， 針對圖像切割的像素級標(biāo)注會多花費十幾倍的時間。因此， 在近些年來半監(jiān)督圖像切割得到了越來越多的關(guān)注。

半監(jiān)督分割的任務(wù)依賴于一部分像素級標(biāo)記圖像和無標(biāo)簽圖像 （通常來說無標(biāo)簽圖像個數(shù)大于等于有標(biāo)簽個數(shù)），其中兩種類型的圖像都遵從相同的數(shù)據(jù)分布。該任務(wù)的挑戰(zhàn)之處在于如何從未標(biāo)記的圖像中提取額外且有用的訓(xùn)練信號，以使模型的訓(xùn)練能夠加強自身的泛化能力。

在當(dāng)前領(lǐng)域內(nèi)有兩個比較火熱的研究方向， 分別是自監(jiān)督訓(xùn)練（self-training） 和 一致性學(xué)習(xí) （consistency learning）。我們的項目主要基于后者來進(jìn)行。

一致性學(xué)習(xí)的介紹

簡單來說， 一致性學(xué)習(xí)（consistency learning）過程可以分為 3 步來描述： 1）。 用不做數(shù)據(jù)增強的 “簡單” 圖像來給像素區(qū)域打上偽標(biāo)簽， 2）。 用數(shù)據(jù)增強 （或擾動） 之后的 “復(fù)雜” 圖片進(jìn)行 2 次預(yù)測， 和 3）。 用偽標(biāo)簽的結(jié)果來懲罰增強之后的結(jié)果。

可是， 為什么要進(jìn)行這 3 步呢？ 先用簡單圖像打標(biāo)簽， 復(fù)雜圖像學(xué)習(xí)的意義在哪？

從細(xì)節(jié)來說， 如上圖所示， 假設(shè)我們有一個像素的分類問題 （在此簡化為 2 分類， 左下的三角和右上的圓圈） 。我們假設(shè)中間虛線為真實分布， 藍(lán)色曲線為模型的判別邊界。

在這個例子中， 假設(shè)這個像素的標(biāo)簽是圓圈， 并且由 1）。 得到的偽標(biāo)簽結(jié)果是正確的 （y_tilde=Circ.）。在 2）。 中如果像素的增強或擾動可以讓預(yù)測成三角類， 那么隨著 3）步驟的懲罰， 模型的判別邊界會 （順著紅色箭頭） 挪向真實分布。由此， 模型的泛化能力得到加強。

由此得出， 在 1）。 中使用 “簡單” 的樣本更容易確保偽標(biāo)簽的正確性， 在 2）。 時使用增強后的 “復(fù)雜” 樣本來確保預(yù)測掉在邊界的另一端來增強泛化能力?？墒窃趯嵺`中，

1）。 沒有經(jīng)受過增強的樣本也很可能被判斷錯 （hard samples）， 導(dǎo)致模型在學(xué)習(xí)過程中打的偽標(biāo)簽正確性下降。

2）。 隨著訓(xùn)練的進(jìn)行， 一般的圖像增強將不能讓模型做出錯誤判斷。這時， 一致性學(xué)習(xí)的效率會大幅度下降。

3）。 被廣泛實用的半監(jiān)督 loss 例如 MSE， 在切割任務(wù)里不能給到足夠的力量來有效的推動判別邊界。而 Cross-entropy 很容易讓模型過擬合錯誤標(biāo)簽， 造成認(rèn)知偏差 （confirmation bias）。

針對這三個問題， 我們提出了：

1）。 新的基于一致性的半監(jiān)督語義分割 MT 模型。通過新引入的 teacher 模型提高未標(biāo)記訓(xùn)練圖像的分割精度。同時， 用置信加權(quán) CE 損失 （Conf-CE） 代替 MT 的 MSE 損失，從而實現(xiàn)更強的收斂性和整體上更好的訓(xùn)練準(zhǔn)確性。

2）。 一種結(jié)合輸入、特征和網(wǎng)絡(luò)擾動結(jié)合的數(shù)據(jù)增強方式，以提高模型的泛化能力。

3）。 一種新型的特征擾動，稱為 T-VAT。它基于 Teacher 模型的預(yù)測結(jié)果生成具有挑戰(zhàn)性的對抗性噪聲進(jìn)一步加強了 student 模型的學(xué)習(xí)效率。

方法介紹

1）。 Dual-Teacher Architecture

我們的方法基于 Mean-Teacher， 其中 student 的模型基于反向傳播做正常訓(xùn)練。在每個 iteration 結(jié)束后， student 模型內(nèi)的參數(shù)以 expotional moving average （EMA）的方式轉(zhuǎn)移給 teacher 模型。

在我們的方法中， 我們使用了兩個 Teacher 模型。在做偽標(biāo)簽時， 我們用兩個 teacher 預(yù)測的結(jié)果做一個 ensemble 來進(jìn)一步增強偽標(biāo)簽的穩(wěn)定性。我們在每一個 epoch 的訓(xùn)練內(nèi)只更新其中一個 teacher 模型的參數(shù)， 來增加兩個 teacher 之間的 diversity。

由于雙 teacher 模型并沒有參加到反向傳播的運算中， 在每個 iteration 內(nèi)他們只會消耗很小的運算成本來更新參數(shù)。

2）。 Semi-supervised Loss

在訓(xùn)練中， teacher 模型的輸出經(jīng)過 softmax 后的置信度代表著它對對應(yīng)偽標(biāo)簽的信心。置信度越高， 說明這個偽標(biāo)簽潛在的準(zhǔn)確率可能會更高。在我們的模型中， 我們首先對同一張圖兩個 teacher 的預(yù)測取平均值。然后通過最后的 confidence 作為權(quán)重， 對 student 模型的輸出做一個基于 cross-entropy 懲罰。同時， 我們會舍棄掉置信度過低的像素標(biāo)簽， 因為他們是噪音的可能性會更大。

3）。 Teacher-based Virtual Adversarial Training （T-VAT）

Virtual Adversarial Training （VAT） 是半監(jiān)督學(xué)習(xí)中常用的添加擾動的方式， 它以部分反向傳播的方式來尋找能最大化預(yù)測和偽標(biāo)簽距離的噪音。

在我們的模型中， dual-teacher 的預(yù)測比學(xué)生的更加準(zhǔn)確， 并且 （由于 EMA 的更新方式使其） 更加穩(wěn)定。我們使用 teacher 模型替代 student 來尋找擾動性最強的對抗性噪音， 進(jìn)而讓 student 的預(yù)測出錯的可能性加大， 最后達(dá)到增強一致性學(xué)習(xí)效率的目的。

4）。 訓(xùn)練流程

i）。 supervised part： 我們用 strong-augmentation 后的圖片通過 cross-entropy 來訓(xùn)練 student 模型。

ii）。 unsupervised part： 我們首先喂給 dual-teacher 模型們一個 weak-augmentation 的圖片， 并且用他們 ensemble 的結(jié)果生成標(biāo)簽。之后我們用 strong-augmentation 后的圖片喂給 student 模型。在通過 encoder 之后， 我們用 dual-teachers 來通過 T-VAT 尋找具有最強擾動性的噪音并且注入到 （student encoded 之后的） 特征圖里， 并讓其 decoder 來做最終預(yù)測。

iii）。 我們通過 dual-teachers 的結(jié)果用 conf-ce 懲罰 student 的預(yù)測

iv）。 基于 student 模型的內(nèi)部參數(shù)， 以 EMA 的方式更新一個 teacher 模型。

實驗

1）。 Compare with SOTAs.

Pascal VOC12 Dataset：

訓(xùn)練 log 可視化鏈接： https://wandb.ai/pyedog1976/PS-MT（VOC12）？workspace=user-pyedog1976

該數(shù)據(jù)集包含超過 13，000 張圖像和 21 個類別。它提供了 1，464 張高質(zhì)量標(biāo)簽的圖像用于訓(xùn)練，1，449 圖像用于驗證，1，456 圖像用于測試。我們 follow 以往的工作， 使了 10582 張低質(zhì)量標(biāo)簽來做擴(kuò)展學(xué)習(xí)， 并且使用了和相同的 label id。

Low-quality Experiments

該實驗從整個數(shù)據(jù)集中隨機(jī) sample 不同 ratio 的樣本來當(dāng)作訓(xùn)練集 （其中包含高質(zhì)量和低質(zhì)量兩種標(biāo)簽）， 旨在測試模型在有不同數(shù)量的標(biāo)簽時所展示的泛化能力。

在此實驗中， 我們使用了 DeeplabV3 + 當(dāng)作架構(gòu)， 并且用 ResNet50 和 ResNet101 得到了所有 ratio 的 SOTA。

High-quality Experiments

該實驗從數(shù)據(jù)集提供的高質(zhì)量標(biāo)簽內(nèi)隨機(jī)挑取不同 ratio 的標(biāo)簽， 來測試模型在極少標(biāo)簽下的泛化能力。我們的模型在不同的架構(gòu)下 （e.g.， Deeplabv3+ and PSPNet） 都取得了最好的結(jié)果。

Cityscapes Dataset

訓(xùn)練 log 可視化鏈接： https://wandb.ai/pyedog1976/PS-MT（City）？workspace=user-pyedog1976

Cityscapes 是城市駕駛場景數(shù)據(jù)集，其中包含 2，975 張訓(xùn)練圖像、500 張驗證圖像和 1，525 張測試圖像。數(shù)據(jù)集中的每張圖像的分辨率為 2，048 ×1，024，總共有 19 個類別。

在 2021 年之前， 大多數(shù)方法用 712x712 作為訓(xùn)練的 resolution， 并且拿 Cross-entropy 當(dāng)作 supervised 的 loss function。在最近， 越來越多的方式傾向于用大 resolution （800x800）當(dāng)作輸入， OHEM 當(dāng)作 supervised loss function。為了公平的對比之前的工作， 我們分別對兩種 setting 做了單獨的訓(xùn)練并且都拿到了 SOTA 的結(jié)果。

2）。 Ablation Learnings.

我們使用 VOC 數(shù)據(jù)集中 1/8 的 ratio 來進(jìn)行消融實驗。原本的 MT 我們依照之前的工作使用了 MSE 的 loss 方式?？梢钥吹剑?conf-CE 帶來了接近 3 個點的巨大提升。在這之后， T-VAT （teacher-based virtual adversarial training）使 student 模型的一致性學(xué)習(xí)更有效率， 它對兩個架構(gòu)帶來了接近 1% 的提升。最后， dual-teacher 的架構(gòu)給兩個 backbone 分別帶來了 0.83% 和 0.84% 的提升。

同時我們對比了多種針對 feature 的擾動的方法， 依次分別為不使用 perturbation， 使用 uniform sample 的噪音， 使用原本的 VAT 和我們提出的 T-VAT。T-VAT 依然帶來了最好的結(jié)果。

3）。 Improvements over Supervised Baseline.

我們的方法相較于相同架構(gòu)但只使用 label part 的數(shù)據(jù)集的結(jié)果有了巨大提升。以 Pascal VOC12 為例， 在 1/16 的比率中 （即 662 張標(biāo)記圖像）， 我們的方法分別 （在 ResNet50 和 ResNet101 中） 超過了基于全監(jiān)督訓(xùn)練的結(jié)果 6.01% 和 5.97%。在其他 ratio 上，我們的方法也顯示出一致的改進(jìn)。

總結(jié)

在本文中，我們提出了一種新的基于一致性的半監(jiān)督語義分割方法。在我們的貢獻(xiàn)中，我們引入了一個新的 MT 模型，它基于多個 teacher 和一個 student 模型，它顯示了對促進(jìn)一致性學(xué)習(xí)的未標(biāo)記圖像更準(zhǔn)確的預(yù)測，使我們能夠使用比原始 MT 的 MSE 更嚴(yán)格的基于置信度的 CE 來增強一致性學(xué)習(xí)的效率。這種更準(zhǔn)確的預(yù)測還使我們能夠使用網(wǎng)絡(luò)、特征和輸入圖像擾動的具有挑戰(zhàn)性的組合，從而顯示出更好的泛化性。

此外，我們提出了一種新的對抗性特征擾動 （T-VAT），進(jìn)一步增強了我們模型的泛化性。