CVPR 2022 | ST++: 半監(jiān)督語義分割中更優(yōu)的自訓練范式

在這里和大家分享一下我們被CVPR 2022錄用的工作"ST++: Make Self-training Work Better for Semi-supervised Semantic Segmentation". 在本工作中，我們在半監(jiān)督語義分割領域重新思考了傳統(tǒng)的多階段自訓練（self-training）范式，并提出兩點關鍵的改進策略，使得這種較為古老的訓練流程仍能達到當前最佳的性能。

作者單位：南京大學、騰訊、東南大學

文章地址：https://arxiv.org/abs/2106.05095

代碼地址：https://github.com/LiheYoung/ST-PlusPlus

Introduction

近來較多的半監(jiān)督學習工作都是基于端到端的框架來做的，學生模型不斷學習教師模型產(chǎn)生的偽標簽。由于模型不斷被更新，偽標簽的質量也會不斷提升，進而持續(xù)促進半監(jiān)督學習的性能。其中比較具有代表性的有半監(jiān)督分類中的FixMatch[1]，半監(jiān)督語義分割中的CutMix-Seg[2]、PseudoSeg[3]、CPS[4]，半監(jiān)督目標檢測中的Unbiased Teacher[5]、Humble Teacher[6]、Soft Teacher[7]等。

盡管這些方法取得了巨大的成功，然而在訓練前期，即使手動設置了置信度閾值，依然存在較多的錯誤偽標簽會誤導學生模型的學習。此外，這種在線學習的范式需要在不同epochs對同一張圖像多次打出偽標簽，比較耗時和耗顯存（一個minibatch中需要同時包含有標簽圖像以及強、弱增廣的無標簽圖像）。因此我們希望重新思考一下傳統(tǒng)的多階段self-training范式在半監(jiān)督語義分割中的應用前景。self-training流程的優(yōu)點在于，模型在有標簽圖像上充分訓練后，對所有無標簽圖像只需要預測一次偽標簽，且偽標簽質量比較有保障。在此基礎上，我們提出了兩種改進方案：

• 在學習無標簽圖像時，在其上施加強數(shù)據(jù)增廣，以增加學習的難度，可以學得額外的信息，并緩解對錯誤偽標簽的過擬合
• 由易至難、從可靠標簽到不可靠標簽，漸進式地利用無標簽圖像及其偽標簽。其中，我們提出基于第一階段訓練過程中偽標簽的穩(wěn)定性來選取可靠的圖像，而非像素。

Background

首先介紹一下最簡單的self-training范式，總共分為三個階段：

1. 【有監(jiān)督預訓練】在有標簽圖像上完全訓練得到一個初始的教師模型
2. 【生成偽標簽】用教師模型在所有的無標簽圖像上預測one-hot偽標簽
3. 【重新訓練】混合有標簽圖像和無標簽圖像及其偽標簽，在其上重新訓練一個學生模型，用于最終的測試

Method

我們的方法分為兩部分，分別是ST和ST++，后者是在前者的基礎上又做了進一步的漸進式選擇策略的改進。

ST: 在重新訓練階段，對無標簽圖像進行強數(shù)據(jù)增廣來學習

我們提出的ST僅需對傳統(tǒng)的self-training范式做很小的改進，就可以顯著提升其性能。

具體來說，由于第二階段預測出的偽標簽仍然是包含較多噪聲的，如果在重新訓練階段直接對這些原圖以及帶噪聲的標簽進行學習，很容易過擬合其中的噪聲標簽；此外，由于學生模型直接學習同樣結構的教師模型產(chǎn)生的偽標簽，此過程并沒有引入額外的信息，學生模型唯一在做的其實只是entropy minimization（因為我們對教師模型的預測結果取了one-hot label）。

基于上述兩點動機，我們提出在重新訓練階段對無標簽圖像進行強數(shù)據(jù)增廣來學習。首先，由于每次無標簽圖像輸入進入模型前都進行了隨機的強數(shù)據(jù)增廣，也就是說盡管是同樣一張圖像，不同epochs見到的輸入也一直在變化，也就沒有固定的輸入-輸出的映射，模型在這種情況下不容易過擬合偽標簽中的噪聲；此外，學生模型是在強增廣的圖像上學習的，可以在教師模型的基礎上學得更加豐富的表征。

這種設計使得學生模型不僅僅是在做entropy minimization，同時由于不同版本的強增廣圖像都受到同樣的偽標簽的監(jiān)督（也就是教師模型生成的固定的one-hot label），此過程也可以看作是在不同的epochs之間對同一張無標簽圖像進行consistency regularization。因此，注入了強增廣操作的self-training范式，也就是我們的ST，同時包含了半監(jiān)督學習中兩種主流的做法，即entropy minimization和consistency regularization。

我們在實驗中采用了四種十分基本的強增廣策略，包括colorjitter、blur、grayscale、以及Cutout。事實上前三種增廣策略的組合就已經(jīng)能夠得到足夠好的結果。

需要注意的是，為了盡可能的減少超參數(shù)，以及增加我們方法的普適性，我們沒有根據(jù)模型預測出的置信度設置一個閾值來選取高質量的偽標簽。并且根據(jù)我們的實驗，這種做法也并沒有在眾多settings上帶來穩(wěn)定的提升。此外，為了盡可能減少訓練時間，我們也沒有進一步的迭代打偽標簽并重訓練（也就是用重新訓練得到的學生模型重新打偽標簽再訓練），但是根據(jù)我們的實驗，這樣做可以進一步提升ST的性能。

下面可以看一下更加詳細的ST偽代碼，流程還是比較直觀的：

ST++: 由易至難、從可靠到不可靠，以圖像級別選取無標簽圖像及其偽標簽

在ST的基礎上，為了進一步緩解錯誤的偽標簽帶來的負面影響，我們提出了ST++，由易至難、從可靠偽標簽到不可靠偽標簽，漸進式地利用無標簽圖像；并且不同于一般做法中選取高置信度的像素，ST++根據(jù)第一階段訓練過程中偽標簽的穩(wěn)定性來選取可靠的圖像。

給定有標簽數(shù)據(jù)集 和無標簽數(shù)據(jù)集 ，ST++首先需要從 中選取出前張最可靠的無標簽圖像及其偽標簽（選取策略后面介紹）構成 ，在 上重新訓練得到一個學生模型 ，此時我們已經(jīng)利用了有標簽圖像和部分較為可靠的圖像及其偽標簽；接著，為了更好地利用剩余的不太可靠的圖像及其偽標簽，即 ，我們利用學得的 對 重新打偽標簽（此處的 比第一階段僅在 上學得的 質量更高），最終在 上重新訓練得到一個最終的學生模型 ，用于測試過程。

下面描述一下我們對可靠的無標簽圖像的選擇策略。我們在實驗中觀察到，比較簡單的圖像在訓練前期就會達到比較高的正確率，且訓練后期偽標簽變化很小；相反，對于比較困難的圖像，模型在訓練的不同epochs預測出的偽標簽往往有較大差異。基于此觀察，我們提出通過度量偽標簽在不同epochs的穩(wěn)定性來確定無標簽圖像及其偽標簽的可靠性。為了這種度量策略更加穩(wěn)定，我們是基于不同epochs的整圖偽標簽之間的meanIOU來算的。具體做法是，在第一階段有標簽圖像上的預訓練過程中保存 個checkpoints，考慮到往往最后一個checkpoint的質量最高，因此對于一張無標簽圖像 ，我們用前 個checkpoints在 上預測出的偽標簽和第 個checkpoint的預測結果算meanIOU，meanIOU越大，說明他們預測出的偽標簽的重合度越高，也就是偽標簽在訓練過程中越穩(wěn)定，其質量也越可靠。形式化描述如下：

其中， 衡量了無標簽圖像 的穩(wěn)定性和可靠性， 表示第 個checkpoint在 上預測出的偽標簽。

獲得了所有無標簽圖像的穩(wěn)定性和可靠性指標 后，我們繼續(xù)基于整圖來進行選取，因為我們認為在語義分割的訓練中，整圖能夠提供比零散的像素更加全局的語義信息。

下面是ST++的偽代碼：

Experiments

Comparison with State-of-the-Art Methods

我們在Pascal VOC 2012和Cityscapes這兩個數(shù)據(jù)集的多種settings上驗證了ST和ST++的有效性。

由于2021年之前的半監(jiān)督語義分割工作大多采用的模型結構是DeepLabv2 with ResNet-101，2021年開始很多工作也增加了PSPNet with ResNet-50，DeepLabv3+ with ResNet-50/101的結果，為了更充分地與更多的工作進行公平對比，我們在Pascal VOC上報告了上述全部的四種模型的結果?？梢钥吹轿覀兊腟T和ST++都取得了state-of-the-art performance。為了更好地展示半監(jiān)督算法的意義，我們也報告了僅利用有標簽圖像的結果，參見每種模型下的第一行SupOnly結果，可以看到半監(jiān)督算法對于SupOnly的提升十分明顯。

Pascal VOC 2012下還有另一種setting，即從原始的高質量標注的訓練集（1464張圖像）中選取有標簽圖像，我們也進行了相應的對比，如下。

我們進一步比較了Cityscapes下的實驗結果。

Ablation Studies

• ST中強數(shù)據(jù)增廣（Strong Data Augmentation, SDA）的意義

在下表中我們首先給出了最原始的self-training的結果（第一行），通過第一行與第三行ST結果的對比，可以看出SDA對于無標簽圖像的作用。為了進一步驗證此提升并非來自于這些增廣策略本身，我們也嘗試將SDA同時加在了有標簽圖像上（第二行），可以看出結果相較于第三行出現(xiàn)了下降。此結果說明SDA的作用并非在于其本身，而是為原本在無標簽圖像上的bootstrapping過程引入了額外的信息，以及緩解了對于噪聲偽標簽的過擬合。

• 不同的SDA的作用

我們整個的SDA包含四種增廣策略：colorjitter、blur、grayscale和Cutout。下圖展示了這四種增廣策略各自的作用，可以看到其中colorjitter的作用相對而言最大，并且前三種樸素的增廣策略的組合就已經(jīng)能夠取得足夠好的結果（下圖棕色柱，73.1）。

• ST++中選取出的可靠樣本和不可靠樣本的偽標簽質量對比

為了更好地展示無標簽圖像中的可靠偽標簽和不可靠偽標簽的質量差距，下圖展示了ST++選取出的可靠/不可靠集合中圖像的mIOU（紫色和棕色）。很容易看出，在不同的settings下，這兩個集合的mIOU差距都大于10%，這說明了首先利用可靠圖像的必要性。在ST++中，我們用在 上訓練得到的學生模型 來給剩下的不可靠集合 重新賦予偽標簽，下圖的Boosted（粉色）一欄展示了在 上， 產(chǎn)生的偽標簽相較于原始的教師模型 產(chǎn)生的偽標簽的提升。結合更新了更好偽標簽的 ，最終我們可以在 上訓練得到更好的學生模型 。

• ST++的提升是否僅僅受益于兩階段的策略？

ST++提出根據(jù)可靠/不可靠集合，把原本的重訓練過程拆成兩個部分。我們需要進一步驗證其提升是否僅僅來自于兩階段的訓練流程，而非對于可靠集合的智能選取。因此，我們隨機從 中選擇50%的圖像作為第一階段的重訓練，訓練完成后再給剩下50%的圖像重新賦予偽標簽，最后用完整的 和 訓練得到最終的模型，我們將其稱為random two-stage re-training。下表展示了random和我們的selective策略的性能差異，可見ST++的提升并不是來自于兩階段的策略。

• ST++中兩階段重訓練的performance

ST++的兩階段重訓練分別是：1）在 上訓練（Re-train #1）；2）給 重新賦予偽標簽后，在 上重新訓練（Re-train #2）。下圖展示了這兩個階段結束后各自的性能。

• ST++中可靠圖像的選取比例

我們默認選取可靠性得分前50%的圖像作為可靠圖像，剩下的作為不可靠圖像。下表展示了不同的選取比例的影響，可以看到ST++對此處的選擇比例比較魯棒。

• ST++中圖像級和像素級選擇策略的比較

在ST++，我們提出基于圖像級別選擇可靠樣本，而非在半監(jiān)督語義分割中通常采用的像素級別樣本[8]。我們認為圖像級別的樣本能夠在語義分割中提供更好的全局語義信息，能夠讓分割模型更加充分的學習。為了驗證這一觀點，我們也比較了圖像級別選擇策略和像素級別選擇策略，下表中展示了兩者性能的差異，可以看到圖像級別的選取穩(wěn)定地優(yōu)于像素級別的選取。

• 進一步進行迭代式重新訓練(iterative re-training)的效果

self-training范式往往可以通過不斷的迭代式的重新訓練獲得進一步的性能提升，具體做法就是用重新訓練得到的學生模型再為無標簽圖像賦予一遍偽標簽，并繼續(xù)重新訓練。我們也嘗試了ST++能否進一步受益于這種迭代式的重新訓練，我們額外增加了一個階段的重訓練（下圖Re-train #3），可以看到ST++仍能獲得進一步的提升。

Conclusion

在本工作中我們調研了傳統(tǒng)的self-training范式在半監(jiān)督語義分割中的應用前景，并提出了兩個關鍵的改進策略，分別是在無標簽圖像上注入強數(shù)據(jù)增廣和基于圖像級別選擇的漸進式重訓練策略。我們發(fā)現(xiàn)，結合了上述改進策略后，傳統(tǒng)的self-traininig范式可以同時享有entropy minimization和consistency regularization的優(yōu)點，并且可以取得優(yōu)于最新的采用端到端訓練的方法的結果。

參考

1. ^FixMatch: Simplifying Semi-Supervised Learning with Consistency and Confidence, NeurIPS 2020
2. ^Semi-supervised semantic segmentation needs strong, varied perturbations, BMVC 2020
3. ^PseudoSeg: Designing Pseudo Labels for Semantic Segmentation, ICLR 2021
4. ^Semi-Supervised Semantic Segmentation with Cross Pseudo Supervision, CVPR 2021
5. ^Unbiased Teacher for Semi-Supervised Object Detection, ICLR 2021
6. ^Humble Teachers Teach Better Students for Semi-Supervised Object Detection, CVPR 2021
7. ^End-to-End Semi-Supervised Object Detection with Soft Teacher, ICCV 2021
8. ^Rethinking Pre-training and Self-training, NeurIPS 2020

公眾號后臺回復“數(shù)據(jù)集”獲取50+深度學習數(shù)據(jù)集下載～