輕量級(jí)MobileSAM：比FastSAM快4倍，處理一張圖像僅需10ms

作者：計(jì)算機(jī)視覺(jué)研究院 2024-06-24 09:35:52

今天分享研究了通用的零樣本分割，在不需要任何訓(xùn)練樣本的情況下，實(shí)現(xiàn)新類(lèi)別的全景、實(shí)例和語(yǔ)義分割。

本文經(jīng)計(jì)算機(jī)視覺(jué)研究院公眾號(hào)授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

論文地址：https://arxiv.org/pdf/2306.14289.pdf

代碼地址：https://github.com/ChaoningZhang/MobileSAM

01 概要簡(jiǎn)介

SAM是一種prompt-guided的視覺(jué)基礎(chǔ)模型，用于從其背景中剪切出感興趣的對(duì)象。自Meta研究團(tuán)隊(duì)發(fā)布SA項(xiàng)目以來(lái)，SAM因其令人印象深刻的零樣本傳輸性能和與其他模型兼容的高度通用性而備受關(guān)注，用于高級(jí)視覺(jué)應(yīng)用，如具有細(xì)粒度控制的圖像編輯。

許多這樣的用例需要在資源受限的邊緣設(shè)備上運(yùn)行，比如移動(dòng)應(yīng)用程序。今天分享中，我們的目標(biāo)是通過(guò)用輕量級(jí)圖像編碼器取代重量級(jí)圖像編碼器，使SAM對(duì)移動(dòng)友好。原始SAM文件中訓(xùn)練這種新SAM的方式會(huì)導(dǎo)致性能不令人滿(mǎn)意，尤其是當(dāng)可用的訓(xùn)練來(lái)源有限時(shí)。

我們發(fā)現(xiàn)，這主要是由圖像編碼器和掩模解碼器的耦合優(yōu)化引起的，因此提出了解耦蒸餾。具體地說(shuō)，將原始SAM中的圖像編碼器ViT-H的知識(shí)提取到一個(gè)輕量級(jí)的圖像編碼器中，該編碼器可以自動(dòng)與原始SAM中的掩碼解碼器兼容。

The overview of Segment Anything Model

Parameters SAM with different image encoders

訓(xùn)練可以在不到一天的時(shí)間內(nèi)在單個(gè)GPU上完成，由此產(chǎn)生的輕量級(jí)SAM被稱(chēng)為MobileSAM，它比原始SAM小60多倍，但性能與原始SAM相當(dāng)。就推理速度而言，MobileSAM每幅圖像運(yùn)行約10ms：圖像編碼器運(yùn)行8ms，掩碼解碼器運(yùn)行2ms。憑借卓越的性能和更高的通用性，我們的MobileSAM比并發(fā)的FastSAM小7倍，快4倍，更適合移動(dòng)應(yīng)用。

02 新框架

Background and Project Goal

Background on SAM

在這里，我們首先總結(jié)SAM的結(jié)構(gòu)及其工作原理。SAM由一個(gè)基于ViT的圖像編碼器和一個(gè)提示引導(dǎo)掩碼解碼器組成。圖像編碼器將圖像作為輸入并生成嵌入，然后將嵌入提供給掩碼解碼器。掩碼解碼器生成一個(gè)掩碼，根據(jù)點(diǎn)（或框）等提示從背景中剪切出任何對(duì)象。此外，SAM允許為同一提示生成多個(gè)掩碼，以解決模糊性問(wèn)題，這提供了寶貴的靈活性。考慮到這一點(diǎn)，這項(xiàng)工作保持了SAM的流水線(xiàn)，首先采用基于ViT的編碼器來(lái)生成圖像嵌入，然后采用提示引導(dǎo)解碼器來(lái)生成所需的掩碼。這條管道是為“分段任何東西”而優(yōu)化設(shè)計(jì)的，可用于“分段所有東西”的下游任務(wù)。

SAM的耦合知識(shí)蒸餾。左圖表示完全耦合蒸餾，右圖表示半耦合蒸餾。

Project goal

該項(xiàng)目的目標(biāo)是生成一個(gè)移動(dòng)友好型SAM（MobileSAM），以輕量級(jí)的方式實(shí)現(xiàn)令人滿(mǎn)意的性能，并且比原始SAM快得多。原始SAM中的提示引導(dǎo)掩碼解碼器的參數(shù)小于4M，因此被認(rèn)為是輕量級(jí)的。給定編碼器處理的圖像嵌入，如他們的公開(kāi)演示中所示，SAM可以在資源受限的設(shè)備中工作，因?yàn)檠诖a解碼器是輕量級(jí)的。然而，原始SAM中的默認(rèn)圖像編碼器是基于ViT-H的，具有超過(guò)600M的參數(shù)，這是非常重量級(jí)的，并使整個(gè)SAM管道與移動(dòng)設(shè)備不兼容。因此，獲得移動(dòng)友好SAM的關(guān)鍵在于用輕量級(jí)的圖像編碼器取代重量級(jí)的圖像編碼器，這也自動(dòng)保持了原始SAM的所有功能和特性。