作者提出了EfficientViT-SAM，這是一系列加速的SAM模型。在保留SAM輕量級的提示編碼器和 Mask 解碼器的同時(shí)，作者用EfficientViT替換了沉重的圖像編碼器。在訓(xùn)練方面，首先從SAM-ViT-H圖像編碼器向EfficientViT進(jìn)行知識蒸餾。隨后，在SA-1B數(shù)據(jù)集上進(jìn)行端到端的訓(xùn)練。得益于EfficientViT的高效性和容量，EfficientViT-SAM在A100 GPU上實(shí)現(xiàn)了48.9的TensorRT速度提升，而且沒有犧牲性能。

代碼和預(yù)訓(xùn)練：https://github.com/mit-han-lab/efficientvit

1 Introduction

Segment Anything Model (SAM)  是一系列在高質(zhì)量數(shù)據(jù)集上預(yù)訓(xùn)練的圖像分割模型，該數(shù)據(jù)集包含1100萬張圖片和10億個(gè) Mask 。SAM 提供了驚人的零樣本圖像分割性能，并在許多應(yīng)用中都有用途，包括增強(qiáng)現(xiàn)實(shí)/虛擬現(xiàn)實(shí)、數(shù)據(jù)標(biāo)注、交互式圖像編輯等。

盡管性能強(qiáng)大，但SAM的計(jì)算量非常大，這在時(shí)間敏感的情境中限制了其適用性。特別是，SAM的主要計(jì)算瓶頸在于其圖像編碼器，在推理時(shí)每張圖像需要2973 GMACs。

為了加速SAM，已經(jīng)進(jìn)行了許多嘗試，用輕量級模型替換SAM的圖像編碼器。例如，MobileSAM 將SAM的ViT-H模型的知識蒸餾到一個(gè)小型視覺 Transformer 中。EdgeSAM 訓(xùn)練了一個(gè)純基于CNN的模型來模仿ViT-H，并采用了一種細(xì)致的蒸餾策略，過程中涉及到提示編碼器和 Mask 解碼器。EfficientSAM 利用MAE預(yù)訓(xùn)練方法來提高性能。

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

作者全面評估了EfficientViT-SAM在一系列零樣本基準(zhǔn)測試上的表現(xiàn)。EfficientViT-SAM在性能和效率上顯著優(yōu)于所有之前的SAM模型。特別是，在COCO數(shù)據(jù)集上，與SAM-ViT-H相比，EfficientViT-SAM在A100 GPU上實(shí)現(xiàn)了48.9倍的吞吐量提升，而mAP沒有下降。

2 Related Work

Segment Anything Model

SAM 在計(jì)算機(jī)視覺領(lǐng)域已經(jīng)獲得廣泛認(rèn)可，作為該領(lǐng)域的一個(gè)里程碑，它在圖像分割方面展示了卓越的性能和泛化能力。SAM 將圖像分割定義為可提示的任務(wù)，旨在給定任何分割提示時(shí)生成有效的分割 Mask 。為了實(shí)現(xiàn)這一目標(biāo)，SAM 使用圖像編碼器和提示編碼器來處理圖像并提供提示。兩個(gè)編碼器的輸出隨后被送入 Mask 解碼器，該解碼器生成最終的 Mask 預(yù)測。

SAM 在一個(gè)大規(guī)模的分割數(shù)據(jù)集上進(jìn)行訓(xùn)練，該數(shù)據(jù)集包含超過1100萬張圖像和超過10億個(gè)高質(zhì)量 Mask ，這使得它能夠在零樣本開放世界分割中表現(xiàn)出強(qiáng)大的能力。SAM 在各種下游應(yīng)用中展示了其高度的適應(yīng)性，包括圖像修復(fù)、目標(biāo)跟蹤和3D生成。然而，SAM的圖像編碼器組件帶來了顯著的計(jì)算成本，導(dǎo)致高延遲，限制了在時(shí)間敏感場景中的實(shí)用性。最近的工作集中在提高SAM的效率，旨在解決其計(jì)算限制。

Efficient Deep Learning Computing

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

3 Method

作者提出了EfficientViT-SAM，該方法利用EfficientViT來加速SAM。特別是，EfficientViT-SAM保留了SAM的提示編碼器和 Mask 解碼器架構(gòu)，同時(shí)用EfficientViT替換了圖像編碼器。作者設(shè)計(jì)了兩系列模型，EfficientViT-SAM-L和EfficientViT-SAM-XL，它們在速度和性能之間提供了平衡。隨后，作者以端到端的方式使用SA-1B數(shù)據(jù)集來訓(xùn)練EfficientViT-SAM。

EfficientViT

EfficientViT 是一系列用于高效高分辨率密集預(yù)測的視覺 Transformer 模型。其核心構(gòu)建模塊是一個(gè)多尺度線性注意力模塊，它通過硬件高效的運(yùn)算實(shí)現(xiàn)了全局感受野和多尺度學(xué)習(xí)。

具體來說，它用輕量級的ReLU線性注意力替代了效率低下的softmax注意力，以擁有全局感受野。通過利用矩陣乘法的結(jié)合性質(zhì)，ReLU線性注意力可以在保持功能的同時(shí)，將計(jì)算復(fù)雜度從二次降低到一次。此外，它還通過卷積增強(qiáng)了ReLU線性注意力，以減輕其在局部特征提取上的局限性。更多細(xì)節(jié)可在原論文中找到。

EfficientViT-SAM

模型架構(gòu)。EfficientViT-SAM-XL的宏觀架構(gòu)如圖2所示。其主干包含五個(gè)階段。類似于EfficientViT，作者在早期階段使用卷積塊，而在最后兩個(gè)階段使用efficientViT模塊。作者通過上采樣和加法融合最后三個(gè)階段的特征。融合后的特征被送入由幾個(gè)融合的MBConv塊組成的 Neck ，然后送入SAM Head 。

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

在端到端的訓(xùn)練階段，作者以相等的概率隨機(jī)選擇框提示和點(diǎn)提示。在點(diǎn)提示的情況下，作者從真實(shí) Mask 中隨機(jī)選擇1-10個(gè)前景點(diǎn)，以確保作者的模型能夠有效應(yīng)對各種點(diǎn)配置。在框提示的情況下，作者使用真實(shí)邊界框。對于EfficientViT-SAM-L/XL模型，作者將最長邊調(diào)整至512/1024，并相應(yīng)地填充較短邊。作者每張圖像選擇多達(dá)64個(gè)隨機(jī)采樣的 Mask 。

為了監(jiān)督訓(xùn)練過程，作者使用Focal Loss和骰子損失的線性組合，F(xiàn)ocal Loss與骰子損失的比例為20:1。類似于SAM中采用的消除歧義的方法，作者同時(shí)預(yù)測三個(gè) Mask ，并且只反向傳播損失最低的那個(gè)。作者還通過添加第四個(gè)輸出Token來支持單一 Mask 的輸出。在訓(xùn)練期間，作者隨機(jī)交替使用兩種預(yù)測模式。

作者使用SA-1B數(shù)據(jù)集對EfficientViT-SAM進(jìn)行了2個(gè)周期的訓(xùn)練，批量大小為256。采用AdamW優(yōu)化器，動(dòng)量參數(shù)設(shè)為0.9，設(shè)為0.999。初始學(xué)習(xí)率對于EfficientViT-SAM-L/XL分別設(shè)定為2e/1e，并使用余弦衰減學(xué)習(xí)率計(jì)劃將其降低至0。在數(shù)據(jù)增強(qiáng)方面，作者應(yīng)用了隨機(jī)水平翻轉(zhuǎn)。

4 Experiment

在本節(jié)中，作者在4.1節(jié)中對EfficientViT-SAM的運(yùn)行時(shí)效率進(jìn)行了全面分析。隨后，作者在COCO 和 LVIS 數(shù)據(jù)集上評估了EfficientViT-SAM的零樣本能力，這些數(shù)據(jù)集在訓(xùn)練過程中未曾遇到。作者執(zhí)行了兩項(xiàng)不同的任務(wù)：4.2節(jié)中的單點(diǎn)有效 Mask 評估以及4.3節(jié)中的邊界框提示實(shí)例分割。這些任務(wù)分別評估了EfficientViT-SAM的點(diǎn)提示和邊界框提示特征的有效性。此外，作者在4.4節(jié)還提供了SGlnW基準(zhǔn)測試的結(jié)果。

Runtime Efficiency

作者比較了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ù)數(shù)量多于其他加速工作，但由于其有效地利用了硬件友好的運(yùn)算符，因此其吞吐量顯著提高。

Zero-Shot Point-Prompted Segmentation

作者在表2中評估了基于點(diǎn)提示對目標(biāo)進(jìn)行分割時(shí)EfficientViT-SAM的零樣本性能。作者采用了文獻(xiàn)[1]中描述的點(diǎn)選擇方法。即初始點(diǎn)被選為距離目標(biāo)邊界最遠(yuǎn)的點(diǎn)。后續(xù)的每個(gè)點(diǎn)都選為距離錯(cuò)誤區(qū)域邊界最遠(yuǎn)的點(diǎn)，該錯(cuò)誤區(qū)域被定義為真實(shí)值和先前預(yù)測之間的區(qū)域。

作者在COCO和LVIS數(shù)據(jù)集上使用1/3/5次點(diǎn)擊報(bào)告性能，以mIoU（平均交并比）作為評價(jià)指標(biāo)。作者的結(jié)果顯示，與SAM相比，性能更優(yōu)，尤其是在提供額外點(diǎn)提示時(shí)。

Zero-Shot Box-Prompted Segmentation

作者評估了EfficientViT-SAM在利用邊界框進(jìn)行目標(biāo)分割中的零樣本性能。首先，作者將真實(shí)邊界框輸入到模型中，結(jié)果展示在表4中。

所有目標(biāo)都報(bào)告了mIoU（平均交并比），并且分別為小型、中型和大型目標(biāo)分別報(bào)告。EfficientViT-SAM在COCO和LVIS數(shù)據(jù)集上顯著超過了SAM。接下來，作者采用一個(gè)目標(biāo)檢測器ViT-Det，并使用其輸出框作為模型的提示。表5的結(jié)果顯示，EfficientViT-SAM相比于SAM取得了更優(yōu)的性能。值得注意的是，即使是EfficientViT-SAM的最輕版本，也顯著優(yōu)于其他加速工作。

另外，作者使用YOLOv8和GroundingDINO 作為目標(biāo)檢測器，在COCO數(shù)據(jù)集上評估了EfficientViT-SAM的性能。YOLOv8是一種實(shí)時(shí)目標(biāo)檢測器，適用于實(shí)際應(yīng)用場景。另一方面，GroundingDINO能夠使用文本提示來檢測目標(biāo)，這使得作者可以基于文本線索進(jìn)行目標(biāo)分割。表6中展示的結(jié)果表明，EfficientViT-SAM相比于SAM具有卓越的性能。

Zero-Shot In-the-Wild Segmentation

野外分割基準(zhǔn)包含25個(gè)零樣本野外分割數(shù)據(jù)集。作者將EfficientViT-SAM與Grounding-DINO結(jié)合，作為框提示，執(zhí)行零樣本分割。每個(gè)數(shù)據(jù)集的全面性能結(jié)果在表3中展示。SAM達(dá)到48.7的mAP，而EfficientViT-SAM獲得了更高的48.9分。

Qualitative Results.

圖3展示了當(dāng)提供點(diǎn)提示、框提示以及SAM模式時(shí)，EfficientViT-SAM的定性分割結(jié)果。結(jié)果顯示，EfficientViT-SAM不僅在分割大型物體上表現(xiàn)出色，也能有效處理小型物體。這些發(fā)現(xiàn)強(qiáng)調(diào)了EfficientViT-SAM卓越的分割能力。

5 Conclusion

在這項(xiàng)工作中，作者引入了EfficientViT-SAM，它使用EfficientViT來替代SAM的圖像編碼器。EfficientViT-SAM在無需犧牲各種零樣本分割任務(wù)性能的情況下，顯著提高了SAM的效率。