漲點(diǎn)！FreeMask：用密集標(biāo)注的合成圖像提升分割模型性能

在這里分享一下我們NeurIPS 2023的工作"FreeMask: Synthetic Images with Dense Annotations Make Stronger Segmentation Models"。在本工作中，我們從語(yǔ)義分割的mask產(chǎn)生大量的合成圖像，并利用這些合成的訓(xùn)練圖像以及他們對(duì)應(yīng)的mask提升在全量真實(shí)數(shù)據(jù)上訓(xùn)練的語(yǔ)義分割模型的性能, e.g., 在ADE20K上，可以將Mask2Former-Swin-T從48.7提升至52.0(+3.3 mIoU)。

代碼：github.com/LiheYoung/FreeMask 論文：https://arxiv.org/abs/2310.15160

在上面的repo中我們也提供了處理過(guò)后的ADE20K-Synthetic數(shù)據(jù)集（包含ADE20K的20倍的訓(xùn)練圖像）和COCO-Synthetic數(shù)據(jù)集（包含COCO-Stuff-164K的6倍的訓(xùn)練圖像），以及結(jié)合合成數(shù)據(jù)訓(xùn)練后更好的Mask2Former、SegFormer、Segmenter模型的checkpoints。

TL;DR

不同于以往的一些工作利用合成數(shù)據(jù)提升few-shot performance（只用少量的真實(shí)數(shù)據(jù)），我們希望利用合成數(shù)據(jù)直接提升fully-supervised performance（用全量的真實(shí)數(shù)據(jù)），這更有挑戰(zhàn)性。

我們利用semantic image synthesis模型來(lái)從semantic mask產(chǎn)生diverse的合成圖像。然而，直接將這些合成圖像加入訓(xùn)練，其實(shí)并不能提升real-image baseline，反而會(huì)損害性能。

因此，我們?cè)O(shè)計(jì)了一個(gè)noise filtering策略以及一個(gè)image re-sampling策略來(lái)更有效地學(xué)習(xí)合成數(shù)據(jù)，最終在ADE20K（20,210張真實(shí)圖像）和COCO-Stuff（164K張真實(shí)圖像）的各種模型上都能取得提升。此外，我們發(fā)現(xiàn)結(jié)合我們的策略后，只利用合成數(shù)據(jù)也可以取得和真實(shí)數(shù)據(jù)comparable的效果。

Take-home Messages

在全量真實(shí)數(shù)據(jù)的基礎(chǔ)上，有效地利用合成數(shù)據(jù)并不容易，需要生成模型足夠好以及設(shè)計(jì)合適的學(xué)習(xí)合成數(shù)據(jù)策略。

在初始階段我們嘗試了多個(gè)GAN-based從mask生成image的模型 (e.g., OASIS[1])，盡管他們的FID指標(biāo)還不錯(cuò)，但遷移到真實(shí)數(shù)據(jù)集上的表現(xiàn)很差（這里的遷移性能，指在合成數(shù)據(jù)集上訓(xùn)練但在真實(shí)驗(yàn)證集上測(cè)試，ADE20K上的mIoU只有～30%）。

基于Stable Diffusion的mask-to-image synthesis model是更好的選擇，如FreestyleNet[2]。

在生成質(zhì)量比較高以及篩選策略比較合理的情況下，joint train合成數(shù)據(jù)和真實(shí)數(shù)據(jù)會(huì)優(yōu)于先用合成數(shù)據(jù)pre-train再用真實(shí)數(shù)據(jù)fine-tune的效果。

Introduction

FreestyleNet基于semantic mask產(chǎn)生的合成圖像，非常diverse以及逼真

Stable Diffusion (SD)等模型已經(jīng)取得了非常好的text-to-image生成效果，過(guò)去一年里，semantic image synthesis領(lǐng)域的工作也開(kāi)始結(jié)合SD的預(yù)訓(xùn)練來(lái)從semantic mask生成對(duì)應(yīng)的image。其中，我們發(fā)現(xiàn)FreestyleNet[2]的生成效果非常好，如上圖所示。因此，我們希望用這些合成圖像以及他們condition on的semantic mask組成新的合成訓(xùn)練樣本對(duì)，加入到原有的真實(shí)訓(xùn)練集中，進(jìn)一步提升模型的性能。

簡(jiǎn)單的失敗嘗試

我們首先檢查了這些合成圖像到真實(shí)圖像的遷移性能，即用合成圖像訓(xùn)練但在真實(shí)圖像的驗(yàn)證集上測(cè)試。我們用SegFormer-B4在真實(shí)圖像上訓(xùn)練可以取得48.5的測(cè)試mIoU，然而用比真實(shí)訓(xùn)練集大20倍的合成數(shù)據(jù)訓(xùn)練后，只得到了43.3 mIoU。此外，我們也嘗試混合真實(shí)數(shù)據(jù)和合成數(shù)據(jù)（會(huì)對(duì)真實(shí)數(shù)據(jù)上采樣到和合成數(shù)據(jù)一樣多，因?yàn)槠滟|(zhì)量更高），然而也只取得了48.2 mIoU，依然落后于僅用真實(shí)圖像訓(xùn)練的結(jié)果。

因此，我們希望能更有效地從這些合成數(shù)據(jù)中進(jìn)行學(xué)習(xí)。

Motivation

由于上述合成數(shù)據(jù)的結(jié)果并不好，我們更仔細(xì)地觀察了一下合成數(shù)據(jù)集，發(fā)現(xiàn)其中存在著很多合成錯(cuò)誤的區(qū)域，如下圖所示的紅色框區(qū)域。這些合成錯(cuò)誤的區(qū)域加入到訓(xùn)練集中后會(huì)嚴(yán)重?fù)p害模型的性能。

紅色框內(nèi)的合成結(jié)果是錯(cuò)誤的

此外，不同的semantic mask對(duì)應(yīng)著不同的場(chǎng)景，不同的場(chǎng)景的學(xué)習(xí)難度其實(shí)是不一樣的，因此它們所需的合成訓(xùn)練圖像的數(shù)量也是不一樣的。如下圖所示，大體上來(lái)看，從左至右semantic mask對(duì)應(yīng)的場(chǎng)景的難度是逐漸增加的，如果對(duì)每張mask產(chǎn)生同樣數(shù)量的合成圖像去學(xué)習(xí)的話，那么這些簡(jiǎn)單的mask對(duì)應(yīng)的圖像就可能會(huì)主導(dǎo)模型的學(xué)習(xí)，模型的學(xué)習(xí)效率就會(huì)很低。

不同的semantic mask對(duì)應(yīng)的場(chǎng)景的難度是不一樣的，大體上來(lái)看，從左至右難度逐漸增加

Method

有了上述的兩個(gè)motivation，具體的做法是非常簡(jiǎn)單的。

Filtering Noisy Synthetic Regions

針對(duì)第一點(diǎn)motivation，我們?cè)O(shè)計(jì)了一個(gè)noise filtering的策略，來(lái)忽略掉合成錯(cuò)誤的區(qū)域。具體來(lái)說(shuō)，我們利用一個(gè)在真實(shí)圖像上訓(xùn)練好的模型去計(jì)算每張合成圖像和它對(duì)應(yīng)的semantic mask之間的pixel-wise loss，直觀來(lái)看，合成錯(cuò)誤的區(qū)域 (pixels)會(huì)呈現(xiàn)比較大的loss。此外，loss的大小也跟不同類別本身的難度有關(guān)。

Hardness-aware Re-sampling

針對(duì)第二點(diǎn)motivation，我們?cè)O(shè)計(jì)了一個(gè)hardness-aware re-sampling策略，來(lái)讓我們的數(shù)據(jù)合成以及訓(xùn)練更加偏向比較難的場(chǎng)景 (semantic mask)，如下圖所示。

為harder的semantic mask產(chǎn)生更多的合成圖像，而減少簡(jiǎn)單的mask的合成圖像

Learning Paradigms

我們探討了兩種從合成圖像中進(jìn)行學(xué)習(xí)的范式，分別是：

Pre-training: 用合成圖像pre-training，然后用真實(shí)圖像進(jìn)一步fine-tuning

Joint training: 混合真實(shí)圖像和合成圖像（會(huì)對(duì)真實(shí)圖像上采樣到與合成圖像同樣的數(shù)量）一起訓(xùn)練

簡(jiǎn)單來(lái)說(shuō)，我們發(fā)現(xiàn)在生成質(zhì)量比較高以及篩選策略比較合理的情況下，joint training的表現(xiàn)會(huì)更好一些。

Experiment

對(duì)比合成圖像和真實(shí)圖像遷移到真實(shí)測(cè)試集的性能

用真實(shí)圖像或合成圖像進(jìn)行訓(xùn)練，并在真實(shí)驗(yàn)證集上測(cè)試

可以看到，在多種模型上，用合成圖像遷移到真實(shí)驗(yàn)證集都可以取得和真實(shí)訓(xùn)練集comparable的效果。

用合成圖像進(jìn)一步提升全監(jiān)督的分割模型性能

Joint training on ADE20K

當(dāng)加入了合成數(shù)據(jù)后，真實(shí)圖像的全監(jiān)督性能獲得了顯著的提升，特別是對(duì)于Mask2Former-Swin-T，我們將mIoU從48.7提升至了52.0(+3.3)；對(duì)于SegFormer-B4，從48.5提升至了50.6 (+2.1)。

Joint training on COCO-Stuff-164K

COCO-Stuff-164K由于原本的真實(shí)數(shù)據(jù)量很大，所以更難提升，但我們?cè)贛ask2Former-Swi-T上仍然取得了+1.9 mIoU的提升。

Pre-training with synthetic images on ADE20K

Ablation Studies

我們的noise filtering和hardness-aware re-sampling的必要性

在沒(méi)有filtering和re-sampling的情況下，F(xiàn)reestyleNet產(chǎn)生的合成圖像在ADE20K和COCO的真實(shí)集上只能得到43.3和48.0的遷移性能，遠(yuǎn)遠(yuǎn)劣于真實(shí)訓(xùn)練圖像的遷移性能(ADE20K: 48.5和COCO: 50.5)，而應(yīng)用我們的策略后，純合成圖像的遷移性能可以提升至48.3 (ADE20K)和49.3 (COCO)，十分接近真實(shí)訓(xùn)練圖像的表現(xiàn)。

在joint training下，我們的兩項(xiàng)策略也是十分有效的，如果沒(méi)有這兩個(gè)策略，混合合成圖像和真實(shí)圖像只能取得48.2的mIoU （真實(shí)圖像：48.5），而加入我們的策略后，可以將真實(shí)圖像48.5的baseline提升至50.6。

合成圖像的數(shù)量

Nmax 控制單張mask最多產(chǎn)生多少?gòu)埡铣蓤D像，在沒(méi)有filtering和re-sampling的情況下，增加合成圖像的數(shù)量反而帶來(lái)了更差的遷移性能；而在經(jīng)過(guò)filtering和re-sampling后，Nmax從6增加到20可以帶來(lái)穩(wěn)定的遷移性能的提升。

更多的ablation studies請(qǐng)參考我們的文章。

Conclusion

在本工作中，我們通過(guò)從semantic mask產(chǎn)生合成圖像，組成大量的合成訓(xùn)練數(shù)據(jù)對(duì)，在ADE20K和COCO-Stuff-164K上顯著提升了多種語(yǔ)義分割模型在全監(jiān)督設(shè)定下的性能。