輕量級MobileSAM:比FastSAM快4倍,處理一張圖像僅需10ms
本文經計算機視覺研究院公眾號授權轉載,轉載請聯系出處。
論文地址:https://arxiv.org/pdf/2306.14289.pdf
代碼地址:https://github.com/ChaoningZhang/MobileSAM
01 概要簡介
SAM是一種prompt-guided的視覺基礎模型,用于從其背景中剪切出感興趣的對象。自Meta研究團隊發(fā)布SA項目以來,SAM因其令人印象深刻的零樣本傳輸性能和與其他模型兼容的高度通用性而備受關注,用于高級視覺應用,如具有細粒度控制的圖像編輯。
許多這樣的用例需要在資源受限的邊緣設備上運行,比如移動應用程序。今天分享中,我們的目標是通過用輕量級圖像編碼器取代重量級圖像編碼器,使SAM對移動友好。原始SAM文件中訓練這種新SAM的方式會導致性能不令人滿意,尤其是當可用的訓練來源有限時。
我們發(fā)現,這主要是由圖像編碼器和掩模解碼器的耦合優(yōu)化引起的,因此提出了解耦蒸餾。具體地說,將原始SAM中的圖像編碼器ViT-H的知識提取到一個輕量級的圖像編碼器中,該編碼器可以自動與原始SAM中的掩碼解碼器兼容。
The overview of Segment Anything Model
Parameters SAM with different image encoders
訓練可以在不到一天的時間內在單個GPU上完成,由此產生的輕量級SAM被稱為MobileSAM,它比原始SAM小60多倍,但性能與原始SAM相當。就推理速度而言,MobileSAM每幅圖像運行約10ms:圖像編碼器運行8ms,掩碼解碼器運行2ms。憑借卓越的性能和更高的通用性,我們的MobileSAM比并發(fā)的FastSAM小7倍,快4倍,更適合移動應用。
02 新框架
Background and Project Goal
- Background on SAM
在這里,我們首先總結SAM的結構及其工作原理。SAM由一個基于ViT的圖像編碼器和一個提示引導掩碼解碼器組成。圖像編碼器將圖像作為輸入并生成嵌入,然后將嵌入提供給掩碼解碼器。掩碼解碼器生成一個掩碼,根據點(或框)等提示從背景中剪切出任何對象。此外,SAM允許為同一提示生成多個掩碼,以解決模糊性問題,這提供了寶貴的靈活性??紤]到這一點,這項工作保持了SAM的流水線,首先采用基于ViT的編碼器來生成圖像嵌入,然后采用提示引導解碼器來生成所需的掩碼。這條管道是為“分段任何東西”而優(yōu)化設計的,可用于“分段所有東西”的下游任務。
SAM的耦合知識蒸餾。左圖表示完全耦合蒸餾,右圖表示半耦合蒸餾。
- Project goal
該項目的目標是生成一個移動友好型SAM(MobileSAM),以輕量級的方式實現令人滿意的性能,并且比原始SAM快得多。原始SAM中的提示引導掩碼解碼器的參數小于4M,因此被認為是輕量級的。給定編碼器處理的圖像嵌入,如他們的公開演示中所示,SAM可以在資源受限的設備中工作,因為掩碼解碼器是輕量級的。然而,原始SAM中的默認圖像編碼器是基于ViT-H的,具有超過600M的參數,這是非常重量級的,并使整個SAM管道與移動設備不兼容。因此,獲得移動友好SAM的關鍵在于用輕量級的圖像編碼器取代重量級的圖像編碼器,這也自動保持了原始SAM的所有功能和特性。
Decoupled distillation for SAM
以ViT-B為圖像編碼器的SAM的耦合蒸餾和解耦蒸餾的比較。與耦合蒸餾相比,解耦蒸餾性能更好,所需計算資源少于1%。
03 實驗
下圖給出了point與bbox提示詞下MobileSAM與原生SAM的結果對比,可以看到:MobileSAM可以取得令人滿意的Mask預測結果。
下圖從Segment everything角度對比了SAM、FastSAM以及MobileSAM三個模型,可以看到:
- MobileSAM與原生SAM結果對齊驚人的好,而FastSAM會生成一些無法滿意的結果
- FastSAM通常生成非平滑的邊緣,而SAM與MobileSAM并沒有該問題