構(gòu)建一個(gè)移動(dòng)端友好的SAM方案MobileSAM

導(dǎo)讀

本文提出一種"解耦蒸餾"方案對(duì)SAM的ViT-H解碼器進(jìn)行蒸餾，同時(shí)所得輕量級(jí)編碼器可與SAM的解碼器"無(wú)縫兼容"。在推理速度方面，MobileSAM處理一張圖像僅需10ms，比FastSAM的處理速度快4倍。

SAM(Segment Anything Model)是一種提示詞引導(dǎo)感興趣目標(biāo)分割的視覺(jué)基礎(chǔ)模型。自提出之日起，SAM引爆了CV社區(qū)，也衍生出了大量相關(guān)的應(yīng)用(如檢測(cè)萬(wàn)物、摳取萬(wàn)物等等)，但是受限于計(jì)算量問(wèn)題，這些應(yīng)用難以用在移動(dòng)端。

本文旨在將SAM的"重量級(jí)"解碼器替換為"輕量級(jí)"以使其可在移動(dòng)端部署應(yīng)用。為達(dá)成該目標(biāo)，本文提出一種"解耦蒸餾"方案對(duì)SAM的ViT-H解碼器進(jìn)行蒸餾，同時(shí)所得輕量級(jí)編碼器可與SAM的解碼器"無(wú)縫兼容" 。此外，所提方案，只需一個(gè)GPU不到一天時(shí)間即可完成訓(xùn)練，比SAM小60倍且性能相當(dāng)，所得模型稱(chēng)之為MobileSAM。在推理速度方面，MobileSAM處理一張圖像僅需10ms(8ms@Encoder，2ms@Decoder)，比FastSAM的處理速度快4倍，這就使得MobileSAM非常適合于移動(dòng)應(yīng)用。

SAM

上圖給出了SAM架構(gòu)示意圖，它包含一個(gè)"重量級(jí)"ViT編碼器與一個(gè)提示詞引導(dǎo)Mask解碼器。解碼器以圖像作為輸入，輸出將被送入Mask解碼器的隱特征(embedding)；Mask解碼器將基于提示詞(如point、bbox)生成用于目標(biāo)分割的Mask。此外，SAM可以對(duì)同一個(gè)提示詞生成多個(gè)Mask以緩解"模棱兩可"問(wèn)題。更多關(guān)于SAM及衍生技術(shù)可參考文末推薦閱讀材料。

延續(xù)SAM架構(gòu)體系：采用輕量級(jí)ViT解碼器生成隱特征，然后采用提示詞引導(dǎo)解碼器生成期望的Mask。本文目標(biāo)：構(gòu)建一個(gè)移動(dòng)端友好的SAM方案MobileSAM，即比原生SAM更快且具有令人滿意的性能?？紤]到SAM不同模塊之間的參數(shù)量問(wèn)題，本文主要聚焦于采用更輕量型的Encoder替換SAM的重量級(jí)Encoder。

實(shí)現(xiàn)方案

Coupled Distillation 一種最直接的方式是參考SAM方案重新訓(xùn)練一個(gè)具有更小Encoder的SAM，見(jiàn)Figure2左圖。如SAM一文所提到：SAM-ViT-H的訓(xùn)練需要256個(gè)A100，且訓(xùn)練時(shí)間達(dá)68小時(shí)；哪怕Encoder為ViT-B也需要128個(gè)GPU。這樣多的資源消耗無(wú)疑阻礙了研究人員進(jìn)行復(fù)現(xiàn)或改進(jìn)。此外，需要注意的是SAM所提供數(shù)據(jù)集的Mask是有預(yù)訓(xùn)練SAM所生成，本質(zhì)上講，重訓(xùn)練過(guò)程也是一種知識(shí)蒸餾過(guò)程，即講ViT-H學(xué)習(xí)到的知識(shí)遷移到輕量級(jí)Encoder中。

Semi-coupled Distillation 當(dāng)對(duì)原生SAM進(jìn)行知識(shí)蒸餾時(shí)，主要困難在于: Encoder與Decoder的耦合優(yōu)化，兩者存在互依賴(lài)。有鑒于此，作者將整個(gè)知識(shí)蒸餾過(guò)程拆解為Encoder蒸餾+Decoder微調(diào)，該方案稱(chēng)之為半耦合蒸餾(Semi-coupled Distillation)，見(jiàn)Figure2右圖。也就是說(shuō)，我們首先對(duì)Encoder進(jìn)行知識(shí)蒸餾，然后再與Decoder進(jìn)行協(xié)同微調(diào)。

Decoupled Distillation 根據(jù)經(jīng)驗(yàn)，我們發(fā)現(xiàn)這種半耦合蒸餾方案仍然極具挑戰(zhàn)性，這是因?yàn)樘崾驹~的選擇具有隨機(jī)性，使得Decoder可變，進(jìn)而導(dǎo)致優(yōu)化變難。有鑒于此，作者提出直接對(duì)原生SAM的編碼器進(jìn)行蒸餾且無(wú)需與Decoder組合，該方案稱(chēng)之為解耦合蒸餾。該方案的一個(gè)優(yōu)勢(shì)在于：僅需使用MSE損失即可，而無(wú)需用于Mask預(yù)測(cè)的Focal與Dice組合損失。

Necessity of Mask Decoder Finetuning 不同于半耦合蒸餾，經(jīng)解耦合蒸餾訓(xùn)練得到的輕量級(jí)Encoder可能與凍結(jié)的Decoder存在不對(duì)齊問(wèn)題。根據(jù)經(jīng)驗(yàn)，我們發(fā)現(xiàn)：該現(xiàn)象并不存在。這是因?yàn)閷W(xué)生Encoder生成的隱特征非常接近于原始老師Encoder生成的隱特征，因此并不需要與Decoder進(jìn)行組合微調(diào)。當(dāng)然，進(jìn)一步的組合微調(diào)可能有助于進(jìn)一步提升性能。

Preliminary Evaluation 上表對(duì)比了耦合蒸餾與解耦合蒸餾的初步對(duì)比?？梢钥吹剑?/p>

從指標(biāo)方面，解耦合蒸餾方案指標(biāo)稍高，0.75mIoU vs 0.72mIoU；

從訓(xùn)練GPU方面，解耦合蒸餾方案僅需兩個(gè)GPU，遠(yuǎn)小于耦合蒸餾方案的128卡，大幅降低了對(duì)GPU的依賴(lài)；

從迭代次數(shù)方面，解耦合蒸餾方案僅需55k次迭代，遠(yuǎn)小于耦合蒸餾方案的180K，大幅降低了訓(xùn)練消耗；

從訓(xùn)練數(shù)據(jù)方面，解耦合蒸餾方案僅需11K數(shù)據(jù)量，遠(yuǎn)小于耦合蒸餾方案的11M，大幅降低了數(shù)據(jù)依賴(lài)。

盡管如此，但ViT-B對(duì)于移動(dòng)端部署仍然非常困難。因此，后續(xù)實(shí)驗(yàn)主要基于TinyViT進(jìn)行。

本文實(shí)驗(yàn)

在具體實(shí)現(xiàn)方面，作者基于ViT-Tiny進(jìn)行本文所提方案的有效性驗(yàn)證，所得MobileSAM與原生SAM的參數(shù)+速度的對(duì)比可參考上表。在訓(xùn)練方面，僅需SA-1B的1%數(shù)據(jù)量+單卡(RTX3090)，合計(jì)訓(xùn)練8個(gè)epoch，僅需不到一天即可完成訓(xùn)練。

上述兩個(gè)圖給出了point與bbox提示詞下MobileSAM與原生SAM的結(jié)果對(duì)比，可以看到：MobileSAM可以取得令人滿意的Mask預(yù)測(cè)結(jié)果。

消融實(shí)驗(yàn)

上表從訓(xùn)練超參bs、epoch、iter等維度進(jìn)行了對(duì)比分析，可以看到：

在同等迭代次數(shù)下，提升bs可以進(jìn)一步提升模型性能；

在同等bs下，提升iter可以進(jìn)一步提升模型性能。

上報(bào)對(duì)比了FastSAM與MobileSAM，可以看到：

從參數(shù)量方面，MobileSAM只有不到10M的參數(shù)量，遠(yuǎn)小于FastSAM的68M；

從處理速度方面，MobileSAM僅需10ms，比FastSAM的40ms快4倍.

上圖從Segment everything角度對(duì)比了SAM、FastSAM以及MobileSAM三個(gè)模型，可以看到：

MobileSAM與原生SAM結(jié)果對(duì)齊驚人的好，而FastSAM會(huì)生成一些無(wú)法滿意的結(jié)果；

FastSAM通常生成非平滑的邊緣，而SAM與MobileSAM并沒(méi)有該問(wèn)題。

最后，補(bǔ)充一下Segment Anything與Segment Everything之間的區(qū)別。

如SAM一文所提到，SAM通過(guò)提示詞進(jìn)行物體分割，也就是說(shuō)，提示詞的作用是指定想分割哪些物體。理論上講，當(dāng)給定合適的提示詞后，任何目標(biāo)都可以被分割，故稱(chēng)之為Segment Anything。

相反，Segment Everything本質(zhì)上是物體候選框生成過(guò)程，不需要提示詞。故它往往被用來(lái)驗(yàn)證下游任務(wù)上的zero-shot遷移能力。

總而言之，Segment Anything解決了任意物體的提示分割基礎(chǔ)任務(wù)；Segment Everything則解決了所有物體面向下游任務(wù)的候選框生成問(wèn)題。