今天給大家分享一個加速視覺分割大模型的工作：EfficientViT-SAM。這是一種新的加速SAM系列。保留了SAM的輕量級提示編碼器和mask解碼器，同時用EfficientViT替換了沉重的圖像編碼器。對于訓(xùn)練，首先從SAM-ViT-H圖像編碼器到EfficientViT的知識蒸餾開始。隨后在SA-1B數(shù)據(jù)集上進行端到端訓(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億個mask。SAM提供了令人驚嘆的zero-shot圖像分割性能，并且具有許多應(yīng)用，包括AR/VR、數(shù)據(jù)標注、交互式圖像編輯等。盡管性能強大，但SAM的計算密集性很高，限制了其在時間敏感場景中的適用性。特別是，SAM的主要計算瓶頸是其圖像編碼器，在推理時每個圖像需要2973 GMACs。

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

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

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

相關(guān)工作

SAM

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

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

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

方法

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

EfficientViT

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

EfficientViT-SAM

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

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

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

在SA-1B數(shù)據(jù)集上對EfficientViT-SAM進行了2個epoch的訓(xùn)練，使用批量大小為256。采用AdamW優(yōu)化器，動量為β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ù)增強，使用隨機水平翻轉(zhuǎn)。

實驗

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

運行時效率

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

zero-shot點提示分割

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

zero-shot框提示分割

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

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

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

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

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

定性結(jié)果

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

結(jié)論

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

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

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