今天給大家分享一個(gè)加速視覺分割大模型的工作：EfficientViT-SAM。這是一種新的加速SAM系列。保留了SAM的輕量級提示編碼器和mask解碼器，同時(shí)用EfficientViT替換了沉重的圖像編碼器。對于訓(xùn)練，首先從SAM-ViT-H圖像編碼器到EfficientViT的知識蒸餾開始。隨后在SA-1B數(shù)據(jù)集上進(jìn)行端到端訓(xùn)練。由于EfficientViT的效率和容量，EfficientViT-SAM在A100 GPU上的TensorRT加速比SAM-ViT-H高出48.9倍，而且不會犧牲性能。

開源地址：https://github.com/mit-han-lab/efficientvit

文章鏈接：https://arxiv.org/pdf/2402.05008

介紹

SAM是一系列在高質(zhì)量數(shù)據(jù)集上預(yù)訓(xùn)練的圖像分割模型，該數(shù)據(jù)集包含1100萬張圖像和10億個(gè)mask。SAM提供了令人驚嘆的zero-shot圖像分割性能，并且具有許多應(yīng)用，包括AR/VR、數(shù)據(jù)標(biāo)注、交互式圖像編輯等。盡管性能強(qiáng)大，但SAM的計(jì)算密集性很高，限制了其在時(shí)間敏感場景中的適用性。特別是，SAM的主要計(jì)算瓶頸是其圖像編碼器，在推理時(shí)每個(gè)圖像需要2973 GMACs。

為了加速SAM，已經(jīng)進(jìn)行了大量努力，將SAM的圖像編碼器替換為輕量級模型。例如，MobileSAM 將SAM的ViT-H模型的知識蒸餾成一個(gè)微型視覺Transformer。EdgeSAM 訓(xùn)練了一個(gè)純CNN模型來模仿ViT-H，采用了一個(gè)精心的蒸餾策略，其中包括提示編碼器和mask解碼器參與過程。EfficientSAM 利用MAE預(yù)訓(xùn)練方法來提高性能。

雖然這些方法可以降低計(jì)算成本，但它們都在性能上存在顯著下降（上圖1）。本文介紹了EfficientViT-SAM來解決這個(gè)限制，通過利用EfficientViT 來替換SAM的圖像編碼器。同時(shí)保留了SAM的輕量級提示編碼器和mask解碼器架構(gòu)。訓(xùn)練過程分為兩個(gè)階段。首先，使用SAM的圖像編碼器作為老師，訓(xùn)練EfficientViT-SAM的圖像編碼器。其次，使用整個(gè)SA-1B數(shù)據(jù)集對EfficientViT-SAM進(jìn)行端到端訓(xùn)練。

在一系列zero-shot基準(zhǔn)測試上對EfficientViT-SAM進(jìn)行了徹底評估。EfficientViT-SAM在所有先前的SAM模型上都提供了顯著的性能/效率提升。特別是，在COCO數(shù)據(jù)集上，與SAM-ViT-H相比，EfficientViT-SAM在A100 GPU上實(shí)現(xiàn)了48.9倍的吞吐量提高，而mAP沒有下降。

相關(guān)工作

SAM

SAM 在計(jì)算機(jī)視覺領(lǐng)域獲得了廣泛認(rèn)可，展示了其在圖像分割中的卓越性能和泛化能力。SAM將圖像分割定義為一個(gè)可提示的任務(wù)，旨在給定任何分割提示生成有效的分割mask。為實(shí)現(xiàn)此目標(biāo)，SAM利用圖像編碼器和提示編碼器處理圖像并提供提示。兩個(gè)編碼器的輸出然后輸入到mask解碼器中，生成最終的mask預(yù)測。SAM在一個(gè)包含超過1100萬張圖像和超過10億個(gè)高質(zhì)量mask的大規(guī)模分割數(shù)據(jù)集上進(jìn)行訓(xùn)練，實(shí)現(xiàn)了穩(wěn)健的zero-shot開放世界分割。SAM在各種下游應(yīng)用中顯示了其高度的多功能性，包括圖像修補(bǔ) 、目標(biāo)跟蹤和3D生成。然而，SAM的圖像編碼器組件會帶來顯著的計(jì)算成本，導(dǎo)致高延遲，限制了其在時(shí)間敏感場景中的實(shí)用性。近期的工作 [2–4, 14] 集中于改善SAM的效率，旨在解決其計(jì)算限制問題。

高效深度學(xué)習(xí)計(jì)算

當(dāng)在邊緣和云平臺上部署深度神經(jīng)網(wǎng)絡(luò)時(shí)，提高其效率至關(guān)重要。我們的工作與高效模型架構(gòu)設(shè)計(jì)相關(guān)，旨在通過將低效的模型架構(gòu)替換為高效的模型架構(gòu)來改善性能和效率之間的平衡。我們的工作還與知識蒸餾相關(guān)，該方法使用預(yù)訓(xùn)練的教師模型來指導(dǎo)學(xué)生模型的訓(xùn)練。此外，可以將EfficientViT-SAM與其他并行技術(shù)結(jié)合使用，進(jìn)一步提高效率，包括剪枝、量化和硬件感知神經(jīng)架構(gòu)搜索。

方法

EfficientViT-SAM，利用EfficientViT 加速SAM。具體來說，我們的方法保留了SAM的提示編碼器和mask解碼器架構(gòu)，同時(shí)將圖像編碼器替換為EfficientViT。設(shè)計(jì)了兩個(gè)系列的模型，EfficientViT-SAM-L 和 EfficientViT-SAM-XL，提供了速度和性能之間的平衡。隨后，使用SA-1B數(shù)據(jù)集對EfficientViT-SAM進(jìn)行端到端訓(xùn)練。

EfficientViT

EfficientViT是一系列用于高效高分辨率密集預(yù)測的視覺Transformer模型。其核心構(gòu)建模塊是多尺度線性注意力模塊，可以通過硬件高效的操作實(shí)現(xiàn)全局感受野和多尺度學(xué)習(xí)。具體來說，它用輕量級的ReLU線性注意力替換了低效的softmax注意力，從而實(shí)現(xiàn)了全局感受野。通過利用矩陣乘法的關(guān)聯(lián)屬性，ReLU線性注意力可以將計(jì)算復(fù)雜度從二次降低到線性，同時(shí)保持功能性。此外，它通過卷積增強(qiáng)了ReLU線性注意力，以減輕其在局部特征提取方面的局限性。

EfficientViT-SAM

「模型架構(gòu)。」 EfficientViT-SAM-XL 的宏觀架構(gòu)如下圖2所示。其骨干網(wǎng)絡(luò)由五個(gè)階段組成。與EfficientViT類似，在前幾個(gè)階段使用卷積塊，而在最后兩個(gè)階段使用EfficientViT模塊。通過上采樣和加法融合來融合最后三個(gè)階段的特征。融合的特征被輸入到由多個(gè)融合MBConv塊組成的neck中，然后傳遞到SAM頭部。

「訓(xùn)練。」 為了初始化圖像編碼器，首先將SAM-ViT-H的圖像嵌入蒸餾到EfficientViT中。采用L2損失作為損失函數(shù)。對于提示編碼器和mask解碼器，通過加載SAM-ViT-H的權(quán)重來初始化它們。然后，以端到端的方式在SA-1B數(shù)據(jù)集上對EfficientViT-SAM進(jìn)行訓(xùn)練。

在端到端訓(xùn)練階段，以相等的概率隨機(jī)選擇box狀提示和點(diǎn)狀提示之間。對于點(diǎn)狀提示，從mask真值中隨機(jī)選擇1-10個(gè)前景點(diǎn)，以確保我們的模型對各種點(diǎn)配置的執(zhí)行效果良好。對于box狀提示，利用真值邊界框。將最長邊縮放到512/1024用于EfficientViT-SAM-L/XL模型，并相應(yīng)地填充較短的邊。每張圖像最多選擇64個(gè)隨機(jī)采樣的mask。為了監(jiān)督訓(xùn)練過程，使用焦點(diǎn)損失和Dice損失的線性組合，焦點(diǎn)損失與Dice損失的比例為20:1。與SAM中采用的方法相似，為了減少歧義，同時(shí)預(yù)測三個(gè)mask，并僅反向傳播最低損失。還通過添加第四個(gè)輸出token來支持單個(gè)mask輸出。在訓(xùn)練過程中，我們隨機(jī)交替兩種預(yù)測模式。

在SA-1B數(shù)據(jù)集上對EfficientViT-SAM進(jìn)行了2個(gè)epoch的訓(xùn)練，使用批量大小為256。采用AdamW優(yōu)化器，動(dòng)量為β1 = 0.9和β2 = 0.999。初始學(xué)習(xí)率設(shè)置為2e?6/1e?6，針對EfficientViT-SAM-L/XL，通過余弦衰減學(xué)習(xí)率調(diào)度衰減到0。關(guān)于數(shù)據(jù)增強(qiáng)，使用隨機(jī)水平翻轉(zhuǎn)。

實(shí)驗(yàn)

在本節(jié)中，首先在第1節(jié)對EfficientViT-SAM的運(yùn)行時(shí)效率進(jìn)行了全面分析。隨后，評估了EfficientViT-SAM在COCO和LVIS數(shù)據(jù)集上的zero-shot能力，這些數(shù)據(jù)集在訓(xùn)練過程中未遇到。執(zhí)行了兩個(gè)不同的任務(wù)：下一個(gè)節(jié)中的單點(diǎn)有效mask評估和再下一節(jié)中的box狀提示實(shí)例分割。這些任務(wù)分別評估了EfficientViT-SAM的點(diǎn)狀提示和盒狀提示特性的有效性。我們還供了SGinW基準(zhǔn)測試的結(jié)果。

運(yùn)行時(shí)效率

我們比較了EfficientViT-SAM與SAM和其他加速工作的模型參數(shù)、MACs和吞吐量。結(jié)果如下表1所示。在單個(gè)NVIDIA A100 GPU上進(jìn)行了TensorRT優(yōu)化的吞吐量測量。結(jié)果顯示，與SAM相比，我們實(shí)現(xiàn)了令人印象深刻的加速，速度提升了17到69倍。此外，盡管EfficientViT-SAM的參數(shù)比其他加速工作多，但由于其有效利用硬件友好的操作符，它表現(xiàn)出了顯著更高的吞吐量。

zero-shot點(diǎn)提示分割

我們評估了EfficientViT-SAM根據(jù)點(diǎn)提示分割對象的zero-shot性能，結(jié)果如下表2所示。我們采用了[1]中描述的點(diǎn)選擇方法。即初始點(diǎn)被選擇為距離對象邊界最遠(yuǎn)的點(diǎn)。每個(gè)后續(xù)點(diǎn)被選擇為距離誤差區(qū)域邊界最遠(yuǎn)的點(diǎn)，誤差區(qū)域被定義為真值和先前預(yù)測之間的區(qū)域。使用COCO和LVIS數(shù)據(jù)集上的1/3/5次點(diǎn)擊進(jìn)行性能報(bào)告，mIoU作為指標(biāo)。結(jié)果表明，與SAM相比，我們在提供額外點(diǎn)提示時(shí)表現(xiàn)出了更優(yōu)異的性能。

zero-shot框提示分割

我們評估了EfficientViT-SAM在使用邊界框進(jìn)行對象分割時(shí)的zero-shot性能。首先將真值邊界框輸入到模型中，并將結(jié)果呈現(xiàn)在下表4中。報(bào)告了所有對象的mIoU，以及分別為小、中、大對象。我們的方法在COCO和LVIS數(shù)據(jù)集上都明顯優(yōu)于SAM。

接下來，我們使用一個(gè)目標(biāo)檢測器ViT-Det，并利用其輸出框作為模型的提示。下表5中的結(jié)果表明，EfficientViT-SAM相比SAM實(shí)現(xiàn)了更優(yōu)異的性能。值得注意的是，即使是EfficientViT-SAM的最輕版本也明顯優(yōu)于其他加速工作的性能。

此外，使用YOLOv8和GroundingDINO作為目標(biāo)檢測器，在COCO數(shù)據(jù)集上評估了EfficientViT-SAM的性能。YOLOv8是一個(gè)適用于實(shí)際應(yīng)用的實(shí)時(shí)目標(biāo)檢測器。另一方面，GroundingDINO能夠使用文本提示檢測目標(biāo)，使能夠根據(jù)文本線索進(jìn)行對象分割。下表6中呈現(xiàn)的結(jié)果顯示了EfficientViT-SAM相對于SAM的出色性能。

zero-shot自然環(huán)境分割

自然環(huán)境分割基準(zhǔn)包括25個(gè)zero-shot自然環(huán)境分割數(shù)據(jù)集。我們將Grounding-DINO作為框提示，為EfficientViT-SAM提供了zero-shot分割。每個(gè)數(shù)據(jù)集的全面性能結(jié)果見下表3。SAM的mAP為48.7，而EfficientViT-SAM的得分更高，為48.9。

定性結(jié)果

下圖3展示了EfficientViT-SAM在提供點(diǎn)提示、框提示和全分割模式下的定性分割結(jié)果。結(jié)果表明，EfficientViT-SAM不僅在分割大目標(biāo)方面表現(xiàn)出色，而且在有效處理小目標(biāo)方面也表現(xiàn)出色。這些發(fā)現(xiàn)突顯了EfficientViT-SAM出色的分割能力。

結(jié)論

EfficientViT-SAM利用EfficientViT替換了SAM的圖像編碼器。EfficientViT-SAM在各種zero-shot分割任務(wù)中實(shí)現(xiàn)了顯著的效率提升，而不會犧牲性能。已將預(yù)訓(xùn)練模型開放源代碼發(fā)布在GitHub上。

本文轉(zhuǎn)自 AI生成未來 ，作者：Zhuoyang Zhang等

原文鏈接:??https://mp.weixin.qq.com/s/NVgNJN7_Tm-1h7zDabt1Dw??