小模型也可以「分割一切」,Meta改進(jìn)SAM,參數(shù)僅為原版5%
對(duì)于 2023 年的計(jì)算機(jī)視覺(jué)領(lǐng)域來(lái)說(shuō),「分割一切」(Segment Anything Model)是備受關(guān)注的一項(xiàng)研究進(jìn)展。
Meta四月份發(fā)布的「分割一切模型(SAM)」效果,它能很好地自動(dòng)分割圖像中的所有內(nèi)容
Segment Anything 的關(guān)鍵特征是基于提示的視覺(jué) Transformer(ViT)模型,該模型是在一個(gè)包含來(lái)自 1100 萬(wàn)張圖像的超過(guò) 10 億個(gè)掩碼的視覺(jué)數(shù)據(jù)集 SA-1B 上訓(xùn)練的,可以分割給定圖像上的任何目標(biāo)。這種能力使得 SAM 成為視覺(jué)領(lǐng)域的基礎(chǔ)模型,并在超出視覺(jué)之外的領(lǐng)域也能產(chǎn)生應(yīng)用價(jià)值。
盡管有上述優(yōu)點(diǎn),但由于 SAM 中的 ViT-H 圖像編碼器有 632M 個(gè)參數(shù)(基于提示的解碼器只需要 387M 個(gè)參數(shù)),因此實(shí)際使用 SAM 執(zhí)行任何分割任務(wù)的計(jì)算和內(nèi)存成本都很高,這對(duì)實(shí)時(shí)應(yīng)用來(lái)說(shuō)具有挑戰(zhàn)性。后續(xù),研究者們也提出了一些改進(jìn)策略:將默認(rèn) ViT-H 圖像編碼器中的知識(shí)提煉到一個(gè)微小的 ViT 圖像編碼器中,或者使用基于 CNN 的實(shí)時(shí)架構(gòu)降低用于 Segment Anything 任務(wù)的計(jì)算成本。
在最近的一項(xiàng)研究中,Meta 研究者提出了另外一種改進(jìn)思路 —— 利用 SAM 的掩碼圖像預(yù)訓(xùn)練 (SAMI)。這是通過(guò)利用 MAE 預(yù)訓(xùn)練方法和 SAM 模型實(shí)現(xiàn)的,以獲得高質(zhì)量的預(yù)訓(xùn)練 ViT 編碼器。
- 論文鏈接:https://arxiv.org/pdf/2312.00863.pdf
- 論文主頁(yè):https://yformer.github.io/efficient-sam/
這一方法降低了 SAM 的復(fù)雜性,同時(shí)能夠保持良好的性能。具體來(lái)說(shuō),SAMI 利用 SAM 編碼器 ViT-H 生成特征嵌入,并用輕量級(jí)編碼器訓(xùn)練掩碼圖像模型,從而從 SAM 的 ViT-H 而不是圖像補(bǔ)丁重建特征,產(chǎn)生的通用 ViT 骨干可用于下游任務(wù),如圖像分類(lèi)、物體檢測(cè)和分割等。然后,研究者利用 SAM 解碼器對(duì)預(yù)訓(xùn)練的輕量級(jí)編碼器進(jìn)行微調(diào),以完成任何分割任務(wù)。
為了評(píng)估該方法,研究者采用了掩碼圖像預(yù)訓(xùn)練的遷移學(xué)習(xí)設(shè)置,即首先在圖像分辨率為 224 × 224 的 ImageNet 上使用重構(gòu)損失對(duì)模型進(jìn)行預(yù)訓(xùn)練,然后使用監(jiān)督數(shù)據(jù)在目標(biāo)任務(wù)上對(duì)模型進(jìn)行微調(diào)。
通過(guò) SAMI 預(yù)訓(xùn)練,可以在 ImageNet-1K 上訓(xùn)練 ViT-Tiny/-Small/-Base 等模型,并提高泛化性能。對(duì)于 ViT-Small 模型,研究者在 ImageNet-1K 上進(jìn)行 100 次微調(diào)后,其 Top-1 準(zhǔn)確率達(dá)到 82.7%,優(yōu)于其他最先進(jìn)的圖像預(yù)訓(xùn)練基線。
研究者在目標(biāo)檢測(cè)、實(shí)例分割和語(yǔ)義分割上對(duì)預(yù)訓(xùn)練模型進(jìn)行了微調(diào)。在所有這些任務(wù)中,本文方法都取得了比其他預(yù)訓(xùn)練基線更好的結(jié)果,更重要的是在小模型上獲得了顯著收益。
論文作者 Yunyang Xiong 表示:本文提出的 EfficientSAM 參數(shù)減少了 20 倍,但運(yùn)行時(shí)間快了 20 倍,只與原始 SAM 模型的差距在 2 個(gè)百分點(diǎn)以內(nèi),大大優(yōu)于 MobileSAM/FastSAM。
在 demo 演示中,點(diǎn)擊圖片中的動(dòng)物,EfficientSAM 就能快速將物體進(jìn)行分割:
EfficientSAM 還能準(zhǔn)確標(biāo)定出圖片中的人:
試玩地址:https://ab348ea7942fe2af48.gradio.live/
方法
EfficientSAM 包含兩個(gè)階段:1)在 ImageNet 上對(duì) SAMI 進(jìn)行預(yù)訓(xùn)練(上);2)在 SA-1B 上微調(diào) SAM(下)。
EfficientSAM 主要包含以下組件:
交叉注意力解碼器:在 SAM 特征的監(jiān)督下,本文觀察到只有掩碼 token 需要通過(guò)解碼器重建,而編碼器的輸出可以在重建過(guò)程中充當(dāng)錨點(diǎn)(anchors)。在交叉注意力解碼器中,查詢來(lái)自于掩碼 token,鍵和值源自編碼器的未掩碼特征和掩碼特征。本文將來(lái)自交叉注意力解碼器掩碼 token 的輸出特征和來(lái)自編碼器的未掩碼 token 的輸出特征進(jìn)行合并,以進(jìn)行 MAE 輸出嵌入。然后,這些組合特征將被重新排序到最終 MAE 輸出的輸入圖像 token 的原始位置。
線性投影頭。研究者通過(guò)編碼器和交叉注意力解碼器獲得的圖像輸出,接下來(lái)將這些特征輸入到一個(gè)小型項(xiàng)目頭(project head)中,以對(duì)齊 SAM 圖像編碼器中的特征。為簡(jiǎn)單起見(jiàn),本文僅使用線性投影頭來(lái)解決 SAM 圖像編碼器和 MAE 輸出之間的特征維度不匹配問(wèn)題。
重建損失。在每次訓(xùn)練迭代中,SAMI 包括來(lái)自 SAM 圖像編碼器的前向特征提取以及 MAE 的前向和反向傳播過(guò)程。來(lái)自 SAM 圖像編碼器和 MAE 線性投影頭的輸出會(huì)進(jìn)行比較,從而計(jì)算重建損失。
經(jīng)過(guò)預(yù)訓(xùn)練,編碼器可以對(duì)各種視覺(jué)任務(wù)的特征表示進(jìn)行提取,而且解碼器也會(huì)被廢棄。特別是,為了構(gòu)建用于分割任何任務(wù)的高效 SAM 模型,本文采用 SAMI 預(yù)訓(xùn)練的輕量級(jí)編碼器(例如 ViT-Tiny 和 ViT-Small)作為 EfficientSAM 的圖像編碼器和 SAM 的默認(rèn)掩碼解碼器,如圖所示 2(底部)。本文在 SA-1B 數(shù)據(jù)集上對(duì) EfficientSAM 模型進(jìn)行微調(diào),以實(shí)現(xiàn)分割任何任務(wù)。
實(shí)驗(yàn)
圖像分類(lèi)。為了評(píng)估本文方法在圖像分類(lèi)任務(wù)上的有效性,研究者將 SAMI 思想應(yīng)用于 ViT 模型,并比較它們?cè)?ImageNet-1K 上的性能。
如表 1 將 SAMI 與 MAE、iBOT、CAE 和 BEiT 等預(yù)訓(xùn)練方法以及 DeiT 和 SSTA 等蒸餾方法進(jìn)行了比較。
SAMI-B 的 top1 準(zhǔn)確率達(dá)到 84.8%,比預(yù)訓(xùn)練基線、MAE、DMAE、iBOT、CAE 和 BEiT 都高。與 DeiT 和 SSTA 等蒸餾方法相比,SAMI 也顯示出較大的改進(jìn)。對(duì)于 ViT-Tiny 和 ViT-Small 等輕量級(jí)模型,SAMI 結(jié)果與 DeiT、SSTA、DMAE 和 MAE 相比有顯著的增益。
目標(biāo)檢測(cè)和實(shí)例分割。本文還將經(jīng)過(guò) SAMI 預(yù)訓(xùn)練的 ViT 主干擴(kuò)展到下游目標(biāo)檢測(cè)和實(shí)例分割任務(wù)上,并將其與在 COCO 數(shù)據(jù)集上經(jīng)過(guò)預(yù)訓(xùn)練的基線進(jìn)行比較。如表 2 所示, SAMI 始終優(yōu)于其他基線的性能。
這些實(shí)驗(yàn)結(jié)果表明,SAMI 在目標(biāo)檢測(cè)和實(shí)例分割任務(wù)中所提供的預(yù)訓(xùn)練檢測(cè)器主干非常有效。
語(yǔ)義分割。本文進(jìn)一步將預(yù)訓(xùn)練主干擴(kuò)展到語(yǔ)義分割任務(wù),以評(píng)估其有效性。結(jié)果如表 3 所示,使用 SAMI 預(yù)訓(xùn)練主干網(wǎng)的 Mask2former 在 ImageNet-1K 上比使用 MAE 預(yù)訓(xùn)練的主干網(wǎng)實(shí)現(xiàn)了更好的 mIoU。這些實(shí)驗(yàn)結(jié)果驗(yàn)證了本文提出的技術(shù)可以很好地泛化到各種下游任務(wù)。
表 4 將 EfficientSAMs 與 SAM、MobileSAM 和 SAM-MAE-Ti 進(jìn)行比較。在 COCO 上,EfficientSAM-Ti 的性能優(yōu)于 MobileSAM。EfficientSAM-Ti 具有 SAMI 預(yù)訓(xùn)練權(quán)重,也比 MAE 預(yù)訓(xùn)練權(quán)重表現(xiàn)更好。
此外, EfficientSAM-S 在 COCO box 僅比 SAM 低 1.5 mIoU,在 LVIS box 上比 SAM 低 3.5 mIoU,參數(shù)減少了 20 倍。本文還發(fā)現(xiàn),與 MobileSAM 和 SAM-MAE-Ti 相比,EfficientSAM 在多次點(diǎn)擊(multiple click)方面也表現(xiàn)出了良好的性能。
表 5 展示了零樣本實(shí)例分割的 AP、APS、APM 和 APL。研究者將 EfficientSAM 與 MobileSAM 和 FastSAM 進(jìn)行了比較,可以看到,與 FastSAM 相比,EfficientSAM-S 在 COCO 上獲得了超過(guò) 6.5 個(gè) AP,在 LVIS 上獲得了 7.8 個(gè) AP。就 EffidientSAM-Ti 而言,仍然遠(yuǎn)遠(yuǎn)優(yōu)于 FastSAM,在 COCO 上為 4.1 個(gè) AP,在 LVIS 上為 5.3 個(gè) AP,而 MobileSAM 在 COCO 上為 3.6 個(gè) AP,在 LVIS 上為 5.5 個(gè) AP。
而且,EfficientSAM 比 FastSAM 輕得多,efficientSAM-Ti 的參數(shù)為 9.8M,而 FastSAM 的參數(shù)為 68M。
圖 3、4、5 提供了一些定性結(jié)果,以便讀者對(duì) EfficientSAMs 的實(shí)例分割能力有一個(gè)補(bǔ)充性了解。
更多研究細(xì)節(jié),可參考原論文。