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