小模型也可以「分割一切」,Meta改進SAM,引數僅為原版5%

 2023-12-07 16:18:43.0

對於 2023 年的計算機視覺領域來說,「分割一切」(Segment Anything Model)是備受關注的一項研究進展。

Meta四月份釋出的「分割一切模型(SAM)」效果,它能很好地自動分割影象中的所有內容

Segment Anything 的關鍵特徵是基於提示的視覺 Transformer(ViT)模型,該模型是在一個包含來自 1100 萬張影象的超過 10 億個掩碼的視覺資料集 SA-1B 上訓練的,可以分割給定影象上的任何目標。這種能力使得 SAM 成為視覺領域的基礎模型,並在超出視覺之外的領域也能產生應用價值。

儘管有上述優點,但由於 SAM 中的 ViT-H 影象編碼器有 632M 個引數(基於提示的解碼器只需要 387M 個引數),因此實際使用 SAM 執行任何分割任務的計算和記憶體成本都很高,這對實時應用來說具有挑戰性。後續,研究者們也提出了一些改進策略:將預設 ViT-H 影象編碼器中的知識提煉到一個微小的 ViT 影象編碼器中,或者使用基於 CNN 的實時架構降低用於 Segment Anything 任務的計算成本。

在最近的一項研究中,Meta 研究者提出了另外一種改進思路 —— 利用 SAM 的掩碼影象預訓練 (SAMI)。這是通過利用 MAE 預訓練方法和 SAM 模型實現的,以獲得高質量的預訓練 ViT 編碼器。

圖片

  • 論文連結:https://arxiv.org/pdf/2312.00863.pdf

  • 論文主頁:https://yformer.github.io/efficient-sam/

這一方法降低了 SAM 的複雜性,同時能夠保持良好的效能。具體來說,SAMI 利用 SAM 編碼器 ViT-H 生成特徵嵌入,並用輕量級編碼器訓練掩碼影象模型,從而從 SAM 的 ViT-H 而不是影象補丁重建特徵,產生的通用 ViT 骨幹可用於下游任務,如影象分類、物體檢測和分割等。然後,研究者利用 SAM 解碼器對預訓練的輕量級編碼器進行微調,以完成任何分割任務。

爲了評估該方法,研究者採用了掩碼影象預訓練的遷移學習設定,即首先在影象解析度為 224 × 224 的 ImageNet 上使用重構損失對模型進行預訓練,然後使用監督資料在目標任務上對模型進行微調。

通過 SAMI 預訓練,可以在 ImageNet-1K 上訓練 ViT-Tiny/-Small/-Base 等模型,並提高泛化效能。對於 ViT-Small 模型,研究者在 ImageNet-1K 上進行 100 次微調後,其 Top-1 準確率達到 82.7%,優於其他最先進的影象預訓練基線。

研究者在目標檢測例項分割語義分割上對預訓練模型進行了微調。在所有這些任務中,本文方法都取得了比其他預訓練基線更好的結果,更重要的是在小模型上獲得了顯著收益。

論文作者 Yunyang Xiong 表示:本文提出的 EfficientSAM 引數減少了 20 倍,但執行時間快了 20 倍,只與原始 SAM 模型的差距在 2 個百分點以內,大大優於 MobileSAM/FastSAM。

圖片

在 demo 演示中,點選圖片中的動物,EfficientSAM 就能快速將物體進行分割:

圖片

EfficientSAM 還能準確標定出圖片中的人:

圖片

試玩地址:https://ab348ea7942fe2af48.gradio.live/

方法

EfficientSAM 包含兩個階段:1)在 ImageNet 上對 SAMI 進行預訓練(上);2)在 SA-1B 上微調 SAM(下)。

圖片

EfficientSAM 主要包含以下元件:

交叉注意力解碼器:在 SAM 特徵的監督下,本文觀察到只有掩碼 token 需要通過解碼器重建,而編碼器的輸出可以在重建過程中充當錨點(anchors)。在交叉注意力解碼器中,查詢來自於掩碼 token,鍵和值源自編碼器的未掩碼特徵和掩碼特徵。本文將來自交叉注意力解碼器掩碼 token 的輸出特徵和來自編碼器的未掩碼 token 的輸出特徵進行合併,以進行 MAE 輸出嵌入。然後,這些組合特徵將被重新排序到最終 MAE 輸出的輸入影象 token 的原始位置。

線性投影頭。研究者通過編碼器和交叉注意力解碼器獲得的影象輸出,接下來將這些特徵輸入到一個小型專案頭(project head)中,以對齊 SAM 影象編碼器中的特徵。為簡單起見,本文僅使用線性投影頭來解決 SAM 影象編碼器和 MAE 輸出之間的特徵維度不匹配問題。

重建損失。在每次訓練迭代中,SAMI 包括來自 SAM 影象編碼器的前向特徵提取以及 MAE 的前向和反向傳播過程。來自 SAM 影象編碼器和 MAE 線性投影頭的輸出會進行比較,從而計算重建損失。

圖片

經過預訓練,編碼器可以對各種視覺任務的特徵表示進行提取,而且解碼器也會被廢棄。特別是,爲了構建用於分割任何任務的高效 SAM 模型,本文采用 SAMI 預訓練的輕量級編碼器(例如 ViT-Tiny 和 ViT-Small)作為 EfficientSAM 的影象編碼器和 SAM 的預設掩碼解碼器,如圖所示 2(底部)。本文在 SA-1B 資料集上對 EfficientSAM 模型進行微調,以實現分割任何任務。

實驗

影象分類。爲了評估本文方法在影象分類任務上的有效性,研究者將 SAMI 思想應用於 ViT 模型,並比較它們在 ImageNet-1K 上的效能。

如表 1 將 SAMI 與 MAE、iBOT、CAE 和 BEiT 等預訓練方法以及 DeiT 和 SSTA 等蒸餾方法進行了比較。

圖片

SAMI-B 的 top1 準確率達到 84.8%,比預訓練基線、MAE、DMAE、iBOT、CAE 和 BEiT 都高。與 DeiT 和 SSTA 等蒸餾方法相比,SAMI 也顯示出較大的改進。對於 ViT-Tiny 和 ViT-Small 等輕量級模型,SAMI 結果與 DeiT、SSTA、DMAE 和 MAE 相比有顯著的增益。

目標檢測例項分割。本文還將經過 SAMI 預訓練的 ViT 主幹擴充套件到下游目標檢測例項分割任務上,並將其與在 COCO 資料集上經過預訓練的基線進行比較。如表 2 所示, SAMI 始終優於其他基線的效能。

圖片

這些實驗結果表明,SAMI 在目標檢測例項分割任務中所提供的預訓練檢測器主幹非常有效。

語義分割。本文進一步將預訓練主幹擴充套件到語義分割任務,以評估其有效性。結果如表 3 所示,使用 SAMI 預訓練主幹網的 Mask2former 在 ImageNet-1K 上比使用 MAE 預訓練的主幹網實現了更好的 mIoU。這些實驗結果驗證了本文提出的技術可以很好地泛化到各種下游任務。

圖片

表 4 將 EfficientSAMs 與 SAM、MobileSAM 和 SAM-MAE-Ti 進行比較。在 COCO 上,EfficientSAM-Ti 的效能優於 MobileSAM。EfficientSAM-Ti 具有 SAMI 預訓練權重,也比 MAE 預訓練權重表現更好。 

此外, EfficientSAM-S 在 COCO box 僅比 SAM 低 1.5 mIoU,在 LVIS box 上比 SAM 低 3.5 mIoU,引數減少了 20 倍。本文還發現,與 MobileSAM 和 SAM-MAE-Ti 相比,EfficientSAM 在多次點選(multiple click)方面也表現出了良好的效能。

圖片

表 5 展示了零樣本例項分割的 AP、APS、APM 和 APL。研究者將 EfficientSAM 與 MobileSAM 和 FastSAM 進行了比較,可以看到,與 FastSAM 相比,EfficientSAM-S 在 COCO 上獲得了超過 6.5 個 AP,在 LVIS 上獲得了 7.8 個 AP。就 EffidientSAM-Ti 而言,仍然遠遠優於 FastSAM,在 COCO 上為 4.1 個 AP,在 LVIS 上為 5.3 個 AP,而 MobileSAM 在 COCO 上為 3.6 個 AP,在 LVIS 上為 5.5 個 AP。

而且,EfficientSAM 比 FastSAM 輕得多,efficientSAM-Ti 的引數為 9.8M,而 FastSAM 的引數為 68M。

圖片

圖 3、4、5 提供了一些定性結果,以便讀者對 EfficientSAMs 的例項分割能力有一個補充性瞭解。

圖片

圖片

圖片

更多研究細節,可參考原論文。

文章來源:機器之心