github
摘要
視頻對象分割和追蹤任意目標領域出現了強大的工具——分割任意模型 2(SAM 2)。SAM 2 實現令人印象深刻的視頻對象分割性能的關鍵組成部分包括用于幀特征提取的大型多階段圖像編碼器,以及存儲過去幀記憶上下文以輔助當前幀分割的記憶機制。然而,多階段圖像編碼器和記憶模塊的高計算復雜度限制了其在實際任務中的應用,例如在移動設備上進行視頻對象分割。為解決這一限制,我們提出了高效追蹤任意目標模型(EfficientTAMs),這是一種輕量級的追蹤任意目標模型,能夠在低延遲和小模型尺寸下產生高質量的結果。
我們的想法是重新審視將普通的非層次化視覺變換器(ViT)作為視頻對象分割的圖像編碼器,并引入一種高效記憶模塊,以降低幀特征提取和當前幀分割記憶計算的復雜性。我們采用普通的輕量級 ViT 和高效記憶模塊構建 EfficientTAMs,并在 SA-1B 和 SA-V 數據集上對模型進行訓練,用于視頻對象分割和追蹤任意目標任務。我們在多個視頻分割基準測試中進行評估,包括半監督視頻對象分割和可提示視頻分割,發現我們提出的采用普通 ViT 的 EfficientTAM 與 SAM 2 模型(HieraB+SAM 2)性能相當,但在 A100 上速度提高了約 2 倍,參數減少了約 2.4 倍。在分割任意圖像任務中,我們的 EfficientTAMs 也優于原始 SAM,A100 上速度提高了約 20 倍,參數減少了約 20 倍。在 iPhone 15 Pro Max 等移動設備上,我們的 EfficientTAMs 能夠以約 10 幀/秒的速度進行視頻對象分割,并保持合理的質量,這突顯了小型模型在設備上視頻對象分割應用中的能力。
1、引言
分割任意模型 2(SAM 2)(Ravi 等人,2024 年)是一種用于統一圖像和視頻對象分割的基礎模型,在多種分割任務中實現了最先進的性能,例如零樣本圖像分割(Kirillov 等人,2023 年;Chen 等人,2023a 年;Deng 等人,2023 年;Chen 等人,2023b 年)、半監督視頻對象分割(Pont-Tuset 等人,2017 年;Xu 等人,2018 年;Oh 等人,2019 年;Bhat 等人,2020 年;Robinson 等人,2020 年;Li 等人,2022b 年;Yang 和 Yang,2022 年;Cheng 和 Schwing,2022 年;Zhang 等人,2023b 年;Wang 等人,2023 年;Wu 等人,2023 年;Cheng 等人,2024 年;Yang 等人,2024 年)、交互式視頻分割(Caelles 等人,2018 年;Heo 等人,2020 年;Cheng 等人,2021a 年;Homayounfar 等人,2021 年;Yang 等人,2023 年;Cheng 等人,2023b 年;Raji? 等人,2023 年;Cheng 等人,2024 年;Delatolas 等人,2024 年)以及其他實際應用(Zhang 等人,2024b 年;Xiong 等人,2024a 年;Shen 等人,2024 年;Zhang 等人,2024a 年;Ding 等人,2024 年;Qiu 等人,2024 年;Tang 等人,2024 年;Zhou 等人,2024 年)。SAM 2 使用多階段圖像編碼器提取層次化幀特征,并引入記憶模塊,以對當前幀特征和觀察到的幀中存儲的記憶進行交叉注意力操作,從而實現在視頻中跨幀一致的對象分割和交互式跟蹤。
盡管具有這些優勢,但 SAM 2 不適合移動部署,尤其是因為大型圖像編碼器(例如 HieraB+)和記憶模塊成本高昂。SAM 2 的默認圖像編碼器 HieraB+(Ryali 等人,2023 年)參數效率低下,例如約 8000 萬個參數。盡管 SAM 2 提供了一個微型版本,但由于層次化圖像編碼器的存在,其運行速度為每秒 43.8 幀,與默認 SAM 2 模型的 47.2 幀相當。此外,記憶標記(例如空間記憶標記和對象指針標記的連接)很長,例如約 3 萬個,這損害了記憶模塊的交叉注意力效率。
在本文中,我們重新審視了用于視頻對象分割和追蹤任意目標的普通非層次化圖像編碼器。我們提出使用輕量級的普通 ViT 圖像編碼器(例如 ViT-Tiny/-Small(Touvron 等人,2021 年))作為 EfficientSAMs(Xiong 等人,2024b 年)所做的一樣,以降低 SAM 2 的復雜性,同時保持良好的性能。此外,我們提出了一種高效的交叉注意力方法,以加速記憶模塊。這是通過利用記憶空間標記的底層結構實現的。我們觀察到記憶空間標記具有很強的局部性,記憶空間標記的更粗略的表示可以很好地替代交叉注意力操作。我們證明,這可以成為原始記憶模塊的良好替代方案。
為了評估我們的方法,我們在視頻和圖像分割基準測試中進行了廣泛的實驗,包括 MOSE、DAVIS、LVOS 和 SA-V 用于視頻分割,以及 SA-23 用于圖像分割。我們的 EfficientTAM 在半監督視頻對象分割方面優于強大的方法,如 Cutie-base、XMem 和 DEVA,同時更加高效。與 SAM 2 相比,我們的 EfficientTAM 性能相當,例如在 SA-V 測試數據集上,分別為 74.5% 和 74.7%,但幀率降低了約 2 倍。在圖像分割基準測試 SA-23 上,我們的 EfficientTAM 實現了 60.7% 的零樣本圖像分割準確率,而 SAM 為 59.1%,SAM 2 為 61.9%。我們還在 iPhone 15 Pro Max 上對我們的 EfficientTAM 模型進行了基準測試,該模型可以以每秒約 10 幀的速度運行,并具有合理的視頻分割性能。
圖1 比較分析。(左)EfficientTAM和SAM 2在單個NVIDIA A100 GPU上的速度比較。雖然SAM 2難以在設備上部署,但我們的EfficientTAM可以在iPhone 15 Pro Max上每幀運行261毫秒。(右)EfficientTAM、SAM 2以及其他高效模型在SA-V測試集上進行零樣本視頻對象分割的FPS/參數/性能比較。我們在單個NVIDIA A100上對所有模型進行了1024×1024輸入分辨率的FPS(每秒幀數)基準測試。
我們的主要貢獻可以總結如下:
- 我們重新審視了使用普通非層次化圖像編碼器 ViT-Tiny/-Small 用于視頻對象分割,并表明普通 ViT 在與 SAM 2 的層次化圖像編碼器相比時可以實現競爭性能。
- 我們提出了一種通過利用底層記憶空間標記結構實現的高效記憶交叉注意力,并展示了其有利的性能。
- 我們提供了 EfficientTAMs,這是一種輕量級的視頻對象分割和追蹤任意目標模型,具有最先進的質量-效率權衡(圖 1),作為 SAM 2 實際部署的補充。
2、相關工作
2.1 視頻對象分割(VOS)
是計算機視覺中的一個基礎任務,它從背景中分割出感興趣的物體,并在視頻中跟蹤目標物體。在無監督設置中(Grundmann 等人,2010 年;Brox 和 Malik,2010 年;Lee 等人,2011 年;Xu 和 Corso,2012 年;Fragkiadaki 等人,2012 年;Perazzi 等人,2012 年;Zhang 等人,2013 年;Li 等人,2013 年;Papazoglou 和 Ferrari,2013 年;Faktor 和 Irani,2014 年;Wang 等人,2015 年;Taylor 等人,2015 年;Perazzi 等人,2016 年),VOS 模型在沒有參考掩碼的情況下分割顯著物體。在半監督設置中(Pont-Tuset 等人,2017 年;Xu 等人,2018 年;Oh 等人,2019 年;Bhat 等人,2020 年;Robinson 等人,2020 年;Li 等人,2022b 年;Yang 和 Yang,2022 年;Cheng 和 Schwing,2022 年;Zhang 等人,2023b 年;Wang 等人,2023 年;Wu 等人,2023 年;Cheng 等人,2024 年;Yang 等人,2024 年),VOS 需要基于目標物體的第一幀掩碼進行跟蹤和分割。對于交互式視頻對象分割(iVOS)(Caelles 等人,2018 年;Heo 等人,2020 年;Cheng 等人,2021a 年;Homayounfar 等人,2021 年;Yang 等人,2023 年;Cheng 等人,2023b 年;Raji? 等人,2023 年;Cheng 等人,2024 年;Delatolas 等人,2024 年),iVOS 模型在用戶引導下(例如點擊、邊界框、涂鴉)對視頻中的物體進行分割(掩碼片段)。在 SAM 2(Ravi 等人,2024 年)中,半監督 VOS 和 iVOS 已被擴展到可提示視覺分割(PVS),模型可以在視頻的任何一幀上通過不同類型的輸入(例如點擊、框、掩碼)進行交互式提示,以分割和跟蹤有效物體。
2.2 視覺變換器(ViTs)
在各種視覺任務中取得了巨大成功,包括圖像分類(Dosovitskiy 等人,2020 年)、目標檢測(Li 等人,2022c 年)、圖像分割(Cheng 等人,2022 年;Kirillov 等人,2023 年)、視頻分類(Fan 等人,2021 年)和視頻對象分割(Duke 等人,2021 年;Yang 等人,2023 年)。最初的 ViT 家族從高效的 ViT-Tiny 擴展到 ViT-Huge,采用普通非層次化架構。也有將變換器與層次化階段結構相結合的層次化視覺變換器,例如 Swin(Liu 等人,2021 年)、MViT(Fan 等人,2021 年;Li 等人,2022d 年)、PViT(Wang 等人,2021 年)和 Hiera(Ryali 等人,2023 年)。盡管取得了成功,但層次化模型通常比普通 ViT 對應模型在實際部署中更慢(Ryali 等人,2023 年)。將 ViT 與卷積(LeCun 等人,1989 年)相結合,探索了快速混合模型,例如 MobileViT(Mehta 和 Rastegari,2021 年)、LeViT(Graham 等人,2021 年)、EfficientFormer(Li 等人,2022e 年)、Next-ViT(Li 等人,2022a 年)、Tiny-ViT(Wu 等人,2022 年)、Castling-ViT(You 等人,2023 年)、EfficientViT(Liu 等人,2023b 年)和 MobileNetv4(Qin 等人,2024 年)。這種朝著構建高效 ViT 的發展與我們構建高效視頻對象分割的 EfficientTAM 工作是正交的。在 SAM(Kirillov 等人,2023 年)和 EfficientSAMs(Xiong 等人,2024b 年)之后,我們正在追求使用普通 ViT 主干網絡進行高效視頻對象分割和追蹤任意目標任務。
2.3 高效注意力機制
該領域已經開發出方法,以減少標準自注意力相對于輸入序列長度的二次成本(Vaswani 等人,2017 年)。局部窗口注意力已應用于 Beltagy 等人(2020 年)和 Zaheer 等人(2020 年),以降低自注意力的復雜性。在 Shen 等人(2018 年)和 Katharopoulos 等人(2020 年)中,提出了線性點積近似方法,通過啟發式地分離鍵和查詢來線性化自注意力中的 softmax 矩陣。在 Choromanski 等人(2020 年)中,Performer 模型使用隨機特征來近似自注意力,實現了線性時間和內存成本。Xiong 等人(2021 年)中的 Nystr?mformer 利用 Nystr?m 方法以線性成本近似自注意力。Wang 等人(2020 年)中的 Linformer 表明自注意力是低秩的,可以通過學習鍵和值的線性投影矩陣來近似。Liu 等人(2023b 年)和 You 等人(2023 年)的方法利用矩陣乘法的結合律,在視覺變換器中實現高效注意力。這一方向已經取得成功,并在視覺任務上取得了不錯的性能。然而,在初步實驗中,我們發現當這些方法被適應性地用于效率提升時,在記憶交叉注意力模塊中的表現不佳。
2.4 分割任意模型
SAM(Kirillov 等人,2023 年)是一種視覺基礎模型,能夠通過交互式提示(例如點和邊界框)分割圖像中的任意對象。SAM 在多種視覺任務中展現出了卓越的零樣本遷移性能以及高度的通用性,這些任務包括廣泛的分割應用(Chen 等人,2023a 年;Cen 等人,2023 年;Deng 等人,2023 年;Chen 等人,2023b 年)、圖像修復(Yu 等人,2023 年)、圖像恢復(Jiang 和 Holz,2023 年)、圖像編輯(Gao 等人,2023 年)、圖像陰影去除(Zhang 等人,2023c 年)、醫學圖像分割(Ma 和 Wang,2023 年)、偽裝目標檢測(Tang 等人,2023 年)、透明目標檢測(Han 等人,2023 年)、基于概念的解釋(Sun 等人,2023 年)、語義通信(Tariq 等人,2023 年)以及目標跟蹤(Cheng 等人,2023b 年;Yang 等人,2023 年)。SAM 在圖像分割任務中表現出的靈活提示的強大能力,激發了將其擴展到視頻對象分割和追蹤任意目標的動機。追蹤任意目標模型(TAM)(Yang 等人,2023 年)結合了 SAM 和 XMem(Cheng 和 Schwing,2022 年),用于交互式視頻對象跟蹤和分割,其中 SAM 負責幀分割,而 XMem 負責跟蹤。SAM-Track(Cheng 等人,2023b 年)通過結合 SAM(Kirillov 等人,2023 年)、DeAOT(Yang 和 Yang,2022 年)以及 Grounding-Dino(Liu 等人,2023a 年)在視頻中執行目標跟蹤和分割。最新的 SAM 2(Ravi 等人,2024 年)通過引入層次化圖像編碼器來處理幀嵌入,并利用記憶模塊將當前幀嵌入基于過往幀進行條件化,從而將 SAM 擴展到視頻分割領域。鑒于移動應用的使用場景以及計算資源受限的應用需求,近期的研究工作致力于降低 SAM 的計算成本,例如 MobileSAM(Zhang 等人,2023a 年)、FastSAM(Zhao 等人,2023 年)和 EfficientSAM(Xiong 等人,2024b 年)。本文則專注于解決 SAM 2 在視頻對象分割和追蹤任意目標的實際部署中的效率挑戰。
3、方法
3.1 預備知識
(1)分割任意目標
SAM(Kirillov 等人,2023 年)包含一個 ViT 圖像編碼器和一個提示引導的掩碼解碼器。編碼器接收圖像并輸出圖像嵌入。然后,解碼器接收圖像嵌入和提示,從而允許從圖像背景中切割出任意目標。SAM 在一個包含超過 10 億個掩碼的圖像數據集上進行訓練。
(2)分割任意目標 2
分割任意目標 2(SAM 2)(Ravi 等人,2024 年)的架構在很大程度上遵循 SAM,由層次化圖像編碼器、提示引導的輕量級掩碼解碼器和一個新的記憶機制組成。SAM 2 使用層次化圖像編碼器 Hiera(Ryali 等人,2023 年)為每一幀生成圖像嵌入。第 3 階段和第 4 階段的步幅為 16 和 32 的特征用于記憶模塊。第 1 階段和第 2 階段的步幅為 4 和 8 的特征雖然不用于記憶模塊,但會被輸入到掩碼解碼器的上采樣層中,用于生成分割掩碼。為了實現穩定的對象跟蹤,SAM 2 采用了一種由輕量級記憶編碼器、輕量級記憶庫和記憶注意力模塊組成的新記憶機制。它存儲來自過往幀的信息,并利用記憶注意力模塊在記憶庫中存儲的記憶和當前幀特征之間執行交叉注意力,從而理解視頻中的時間依賴性。
(3)記憶注意力模塊由一系列變換器塊組成
每個塊包含自注意力、交叉注意力和 MLP。第一個變換器塊以當前幀的圖像嵌入作為輸入。每個變換器塊的核心組件,即交叉注意力,將當前幀嵌入和存儲在記憶庫中的記憶結合起來,生成包含時間對應信息的嵌入。對于記憶標記,它包括兩部分:來自記憶編碼器的空間嵌入標記和來自掩碼解碼器的對象級指針標記。假設空間標記的數量為 n,對象級指針標記的數量為 P,通道維度為 dm,記憶標記可以表示為
設 L?為標記數量, 為自注意力后輸入幀特征的每個標記的維度,
。輸入序列
?被線性投影到輸入查詢
,記憶標記
?被線性投影到鍵
?和值
,其中 d 是查詢、鍵和值的嵌入維度。應用于查詢 Q 、鍵 K?和值 V 的縮放點積交叉注意力機制可以正式寫為:
? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ? ??
其中,softmax 操作是按行應用的。記憶模塊中使用的是單頭交叉注意力。在后續討論中,我們也將鍵和值視為記憶標記以簡化問題。
(4)效率瓶頸
盡管層次化圖像編碼器在多尺度幀特征提取以及交叉注意力在整合當前幀特征與存儲記憶方面具有優勢,但它們給 SAM 2 的實際部署帶來了挑戰。由于圖像編碼器的層次化設計以及層次化特征的使用,效率低下的 SAM 2(tiny)版本甚至顯示出與基礎 SAM 2 相當的每秒幀數(FPS),分別為 47.2 FPS 和 43.8 FPS,這也使得 SAM 2 很難在移動設備上部署。此外,用于記憶模塊中執行交叉注意力的鍵和值中的標記數量非常多,例如達到 30,000 個。這導致在執行交叉注意力時計算和內存成本很高,從而成為記憶模塊在實際部署中的效率瓶頸。
3.2?高效視頻對象分割和追蹤任意目標
我們現在著手解決 SAM 2 的效率問題,以構建高效的視頻對象分割和追蹤任意目標模型,即 EfficientTAM。受到 SAM 和 EfficientSAM 在高質量分割性能方面的啟發,我們重新審視使用普通、非層次化的輕量級圖像編碼器(如 ViT-Small/ViT-Tiny)進行幀特征提取。我們發現,使用普通的 ViT 進行幀特征提取使 EfficientTAM 高度高效,并且能夠在移動設備上部署。此外,我們引入了一個高效的記憶模塊,通過提出一種高效的交叉注意力操作來減少計算和內存成本。基于這兩個設計,我們通過在很大程度上遵循 SAM 2 來構建高效的視頻對象分割和追蹤任意目標模型。圖 2 展示了我們提出的 EfficientTAM 的概述。
圖2 EfficientTAM架構。我們提出的EfficientTAM采用普通輕量級ViT圖像編碼器進行幀特征提取。提出了一種高效的記憶交叉注意力機制,通過利用記憶空間嵌入的強局部性,進一步提高EfficientTAM的效率。EfficientTAM在SA-1B(圖像)和SA-V(視頻)上進行了全面訓練,以實現統一的圖像和視頻分割。
(1)高效圖像編碼器
圖像編碼器的作用是為每個高分辨率幀生成特征嵌入。我們使用一個經過 SAMI(Xiong 等人,2024b)預訓練的普通 ViT 圖像編碼器(Dosovitskiy 等人,2020 年;Touvron 等人,2021 年)來提取幀特征。與 SAM 2 的圖像編碼器不同,我們的圖像編碼器僅提供單尺度特征圖,并且在解碼過程中,掩碼解碼器中沒有其他特征被添加到上采樣層中以生成分割掩碼。我們采用具有 16×16 塊大小的輕量級圖像編碼器 ViT-Small 和 ViT-Tiny。按照(Li 等人,2022c)的方法,我們使用 14×14 的非重疊窗口注意力和 4 個等間距的全局注意力塊,高效地從高分辨率幀中提取特征。我們的圖像編碼器輸出一個單尺度特征嵌入,其分辨率降低了 16 倍,它接收高分辨率(例如 1024×1024)幀作為輸入,并將其轉換為尺寸縮小后的密集嵌入,大小為 64×64。
(2)高效記憶模塊
記憶模塊利用之前幀的信息來促進一致的對象跟蹤。由于其記憶標記序列較長,SAM 2(Ravi 等人,2024 年)中記憶模塊的交叉注意力是主要的效率瓶頸。我們現在討論如何利用記憶標記的底層結構——空間記憶標記內的局部平滑性(強局部性)——來產生一個有效的替代方案。
考慮兩個連續的空間記憶標記,ki? 和 ki+1?,局部平滑性意味著 ∥ki??ki+1?∥22?≤n2cK??,對于 i=1,…,n?1,其中 cK? 是一個正常數。這表明,給定一個足夠小的局部窗口 lw?×lh?,使用單個標記來代表同質窗口中的其他標記,可能提供完整空間記憶標記集 Ks? 的粗略表示 K~s?。我們可以通過在每個窗口 lw?×lh? 中重復單個標記 lw×lh 次,從 K~s? 構造出與 Ks? 大小相同的良好替代品 Kˉs?。在平滑性假設下,Kˉs? 與 Ks? 不會相差太遠。經驗上,我們觀察到空間記憶標記的粗略表示是完整空間記憶標記的良好替代品。圖 3 確認了輸入鍵和值的粗略表示與記憶模塊中交叉注意力的原始鍵和值接近。
圖3 展示記憶模塊中交叉注意力的鍵和值的強局部性的一個例子。鍵和值是一個28700×256的矩陣。交叉注意力是一個4096×256的矩陣。為了簡化可視化和比較,我們只繪制了大小為320×256的頂部矩陣。我們使用一個平均化的單個標記來代表2×2大小的同質窗口中的其他標記,以獲得鍵和值的粗略表示。在右側,我們可視化了方程(1)中的原始交叉注意力和方程(5)中的高效交叉注意力之間的差異;在Frobenius范數下,相對于原始交叉注意力的相對誤差為0.03。
利用交叉注意力中高度相關的鄰近標記,我們執行平均池化,以高效地為鍵 K 和值 V 計算粗略表示。對于輸入空間標記 Ks?=[k11?,…,k1h?;…;kw1?,…,kwh?],其中 w×h 是分辨率大小,我們將 n=w×h 個標記劃分為 k=w~×h~ 個矩形池化區域,并計算每個區域的平均標記。為了簡化,我們假設 w 能被 w~ 整除,h 能被 h~ 整除。記 lw?=w~w?,lh?=h~h?。可以通過對每個區域進行平均來計算 K~s? 和 V~s?:
其中 ,
。這種標記池化方案只需要對標記進行一次掃描,從而高效地生成粗略標記。我們發現,使用 2×2 的窗口大小進行平均池化足以確保對空間記憶標記的良好近似。
假設 ?是記憶空間鍵?
?的粗略表示,我們可以從
?構造一個與
?大小相同的良好替代品
,通過將每個?
?,重復
?次,可以寫為:
同樣地,我們將每個 v~i?,i=1,…,w~h~,重復 lw?×lh? 次來構造 Vˉs?∈Rn×d 作為值 Vs?∈Rn×d 的良好替代品,可以寫為:
然后我們將這些粗略的空間標記與對象指針標記連接起來,構造? 和
,作為原始記憶標記 K 和 V 的良好替代品。對于粗略的記憶標記 Kˉ 和 Vˉ,我們有:
其中 。我們在附錄中提供了方程 (7) 的證明。由于 Kˉ 和 Vˉ 分別是 K 和 V 的良好替代品,我們得到了原始交叉注意力 softmax
在方程 (1) 中的良好替代品:
有了方程(7),我們得到了一個高效的交叉注意力版本
(3)鏈接到高效的交叉注意力變體
有趣的是,我們可以基于我們在方程(5)中提出的高效交叉注意力找到一些交叉注意力變體。我們注意到有一個常數用于平衡粗略空間標記和對象指針標記之間的注意力分數,避免在池化后減少對空間標記的注意力。如果我們移除這個常數,它可能會導致一個使用平均池化來替代可學習投影的Linformer變體。而不是移除這個常數,我們在方程(6)中將其添加到鍵中,以調節粗略空間標記和對象指針標記之間的注意力,從而獲得另一個變體。
其中:
由于空間記憶嵌入具有很強的局部性,因此實現原始交叉注意力的良好替代是可行的。我們的高效交叉注意力與原始交叉注意力非常接近,如圖3所示。
4、實驗
4.1 實驗設置
(1)預訓練
SA-1B數據集包含1100萬張多樣化、高分辨率的圖像,以及11億個高質量的分割掩碼。與(Ravi等人,2024年)類似,我們在不含記憶組件的情況下,在SA-1B數據集(Kirillov等人,2023年)上對EfficientTAM進行了90k步的預訓練。我們的ViT圖像編碼器是從預訓練的ViTs(Xiong等人,2024b)初始化的。我們使用了AdamW優化器(Loshchilov和Hutter,2019年),其動量參數為(β1=0.9,β2=0.999),全局批量大小為256,初始學習率為4e-4。學習率通過反平方根學習率調度(Zhai等人,2022年)進行衰減,包含1k次迭代的線性熱身和5k次迭代的線性冷卻。我們將權重衰減設置為0.1。我們沒有對圖像編碼器應用Drop Path。層間衰減(Clark等人,2020年)設置為0.8。我們應用水平翻轉增強,并將輸入圖像分辨率調整為1024×1024。我們將每張圖像的訓練限制為64個掩碼。我們的模型在256個A100 GPU上進行預訓練,每個GPU擁有80GB的顯存,使用焦點損失和Dice損失的線性組合進行掩碼預測(例如,比例為20:1)。在訓練過程中使用了Bfloat16。
(2)完整訓練數據集
按照(Ravi等人,2024年)的方法,我們在SA-V數據集(Ravi等人,2024年)和SA-1B(Kirillov等人,2023年)的10%子集上訓練包含記憶組件的EfficientTAM。SA-V是一個大規模且多樣化的視頻分割數據集,包含51,000個視頻,這些視頻跨越47個國家,擁有600,000個掩碼注釋,涵蓋了整個對象及其部分。SA-V視頻的分辨率從240p到4K不等,時長從4秒到138秒不等。與SAM 2不同,我們在訓練過程中沒有使用其他開源數據集或內部數據集,以確保與基線的公平比較。
(3)完整訓練實現細節
與(Ravi等人,2024年)類似,在預訓練后,我們對EfficientTAM進行了300k步的訓練。我們使用了AdamW優化器(Loshchilov和Hutter,2019年),其動量參數為(β1=0.9,β2=0.999),批量大小為256,圖像編碼器的初始學習率為6e-5,模型其他組件的初始學習率為3e-4。學習率通過余弦調度進行衰減,并包含15k次迭代的線性熱身。我們將權重衰減設置為0.1。我們沒有對圖像編碼器應用Drop Path。層間衰減(Clark等人,2020年)設置為0.8。我們應用水平翻轉圖像增強,并將輸入圖像分辨率調整為1024×1024。對于視頻,我們應用水平翻轉增強,25度的仿射變換和20度的剪切,亮度為0.1、對比度為0.03、飽和度為0.03的顏色抖動,以及概率為0.05的灰度增強。我們將每張圖像的訓練限制為64個掩碼,每幀視頻的訓練限制為3個掩碼。我們的模型在256個A100-80G GPU上進行訓練,使用焦點損失和Dice損失的線性組合進行掩碼預測,平均絕對誤差損失進行IoU預測,以及交叉熵損失進行對象預測。線性組合損失的比例為20:1:1:1。在訓練過程中使用了Bfloat16。
(4)下游任務/數據集/模型
我們考慮了零樣本視頻任務,包括可提示視頻分割和半監督視頻對象分割,以及零樣本圖像任務,以展示EfficientTAM在圖像和視頻分割方面的競爭力。對于零樣本圖像任務,我們在37個數據集上評估EfficientTAM,其中包括23個SA-23(Kirillov等人,2023年)數據集和14個在(Ravi等人,2024年)中介紹的視頻數據集。對于零樣本視頻任務,我們在9個密集注釋的數據集上評估EfficientTAM的可提示視頻分割性能。我們使用17個視頻數據集,在不同的提示下評估交互式半監督VOS設置下的零樣本準確率。對于提供第一幀真實掩碼的標準半監督VOS設置,我們使用MOSE(Ding等人,2023年)、DAVIS2017(Pont-Tuset等人,2017年)、LVOS(Hong等人,2024年)、SA-V(Ravi等人,2024年)和YTVOS(Xu等人,2018年)來衡量VOS的準確率。我們建議讀者參考(Kirillov等人,2023年;Ravi等人,2024年)以了解這些數據集的詳細信息。模型方面,我們使用EfficientTAM進行零樣本圖像和視頻任務。
(5)基線和評估指標
對于提供第一幀掩碼的標準半監督VOS任務,我們將EfficientTAM的性能與SAM 2(Ravi等人,2024年)、Cutie-base(Cheng等人,2024年)、DEVA(Cheng等人,2023a年)、XMem(Cheng和Schwing,2022年)等進行了比較。對于使用不同提示的零樣本可提示視頻分割任務和交互式半監督視頻對象分割任務,我們將我們的方法與SAM2(Ravi等人,2024年)、SAM+XMem++(Ravi等人,2024年)和SAM+Cutie(Ravi等人,2024年)進行了比較。對于零樣本圖像分割任務,我們與SAM(Kirillov等人,2023年)和SAM2(Ravi等人,2024年)進行了比較。請注意,我們使用的是開源版本的SAM 2(未在MOSE/LVOS/YTVOS上進行訓練)進行比較。我們還承認了最近發布的經過長時間記憶上下文訓練的SAM 2.1。
我們使用聯合J(區域相似性)&F(輪廓準確性)的準確率指標來評估零樣本視頻分割任務;使用mIoU(平均交并比)來評估零樣本圖像分割任務。對于效率指標,我們比較了模型參數的數量、GPU(例如A100)上的推理吞吐量以及移動設備(例如iPhone 15 Pro Max)上的延遲。我們遵循SAM 2(Ravi等人,2024年)來報告指標。在MOSE、LVOS和YTVOS上提供主要結果時,我們向它們的基準測試服務器提交以在MOSE val、LVOS val和YTVOS2019 val上進行最終性能評估。對于消融研究,我們在MOSE開發集上進行評估,該集包含從MOSE訓練分割中隨機抽取的200個視頻(Ravi等人,2024年)。
4.2 主要結果
(1)標準半監督視頻對象分割
半監督視頻對象分割是基于第一幀的真實掩碼在視頻中進行對象分割和跟蹤的過程。我們遵循SAM 2(Ravi等人,2024年)的方法,并報告了我們在這一標準半監督視頻對象分割任務上的準確率。我們還在單個A100 GPU上以批量大小為1報告了延遲。我們使用不同的圖像編碼器(ViT-Tiny和ViT-Small)和記憶模塊(原始記憶塊和具有2×2窗口池化的高效記憶塊)對EfficientTAM進行了評估,以在效率和準確率之間進行權衡。EfficientTAM-S表示使用ViT-Small圖像編碼器和原始記憶塊的EfficientTAM,而EfficientTAM-S/2表示使用ViT-Small圖像編碼器和具有2×2窗口池化的高效記憶塊的EfficientTAM。表1對比了我們的EfficientTAM與VOS基線,包括SAM 2(Ravi等人,2024年)、Cutie-base(Cheng等人,2024年)和XMem(Cheng和Schwing,2022年)。在SA-V測試集上,我們的EfficientTAM-S實現了74.5 J&F,分別比Cutie-base、Cutie-base+和XMem高出12.2、12.9和14.4。在長期視頻對象分割基準測試LVOS上,我們也可以看到我們的EfficientTAM-S大幅超越了Cutie-base和XMem。請注意,我們的EfficientTAM-S在5個視頻基準測試中僅比SAM 2低<2 J&F或G,但速度提高了約2倍,參數減少了約2.4倍。此外,具有高效記憶注意力的EfficientTAM略微遜色于具有原始記憶注意力的EfficientTAM,但在移動設備上的速度提升顯著,iPhone 15上的延遲減少了>2倍。例如,EfficientSAM-S在SA-V測試集上實現了74.5 J&F,每幀運行時間為1010.8毫秒。具有高效交叉記憶注意力的EfficientSAM-S/2實現了74.0 J&F,僅需450毫秒。這些結果展示了EfficientTAM在半監督視頻對象分割中的卓越優勢,并驗證了我們的方法在實際部署中的優勢。
(2)可提示視頻分割
與SAM 2(Ravi等人,2024年)類似,我們使用兩種設置來評估可提示視頻分割:離線評估和在線評估。對于離線評估,我們多次遍歷視頻,以標記與最大模型誤差相關的幀。對于在線評估,我們只遍歷一次視頻來標記幀。在包括EndoVis、ESD、LVOSv2、LV-VIS、UVO、VOST、PUMaVOS、Virtual KITTI 2和VIPSeg在內的9個密集注釋的視頻數據集上,每幀使用3次點擊進行評估。報告了在1到8個交互幀上的平均J&F準確率。圖4展示了我們的方法與包括SAM 2、SAM+XMem++和SAM+Cutie在內的強大基線之間的比較。EfficientTAM在兩種評估設置中均優于SAM+XMem++和SAM+Cutie。EfficientTAM還縮小了SAM 2在離線和在線設置之間的差距。具體來說,在8個帶有3次點擊的注釋幀中,EfficientTAM-S和EfficientTAM-S/2在離線評估設置中平均實現了約82 J&F,在線評估中平均實現了約81 J&F,比SAM+XMem++和SAM+Cutie高出>3 J&F,并縮小了SAM 2的差距。這組實驗進一步驗證了EfficientTAM在可提示視頻分割方面的有效性。
(3)交互式半監督視頻對象分割
我們還遵循 SAM 2 的方法,僅在第一幀提供點擊、框或掩碼提示,對我們的方法在交互式半監督視頻對象分割任務上進行了評估。在表 2 中,我們針對每種類型的提示,報告了在 17 個視頻數據集上的平均 J&F 準確率。我們觀察到,EfficientTAM 在不同輸入提示下均優于 SAM+XMem++ 和 SAM+Cutie。我們還注意到 EfficientTAM 與 SAM 2 之間的差距縮小了。僅使用 1 次點擊時,我們的 EfficientTAM-S 獲得了 63 J&F 準確率,比 SAM+XMem++ 和 SAM+Cutie 高出 6 J&F,與 SAM 2 相比僅有輕微的損失,即 1.3 J&F。總體而言,EfficientTAM 在使用不同提示的交互式半監督 VOS 任務上表現良好。
(4)圖像上的分割任意目標
我們現在對我們的模型在圖像上的分割任意目標任務進行評估。在表 3 中,我們在 SA-23 基準測試以及 SAM 2(Ravi 等人,2024 年)中引入的 14 個視頻數據集的新基準測試中,報告了 1 次點擊和 5 次點擊的 mIoU 準確率。我們將我們的 EfficientTAM 與 SAM(ViT-H)和 HQ-SAM(ViT-H)進行了比較。我們的 EfficientTAM-S 在 1 次點擊準確率上比 SAM(ViT-H)提高了 2.6 mIoU,比 HQ-SAM(ViT-H)提高了 1.6 mIoU。對于 5 次點擊,我們觀察到與 SAM(ViT-H)和 HQ-SAM(ViT-H)相比有持續的改進。我們還在 SA-23 的視頻基準測試以及包含 14 個新視頻的基準測試中注意到顯著的改進。這表明我們的 EfficientTAM 在圖像和視頻分割方面都很強大。
(5)定性評估
圖 5 展示了兩個視頻示例。我們比較了在第一幀提示掩碼的情況下 EfficientTAM 和 SAM 2 的表現。我們發現,我們的 EfficientTAM 能夠像 SAM 2 一樣為目標對象生成高質量的掩碼片段。更多視頻示例見附錄。這些結果表明,我們的 EfficientTAM 具有與 SAM 2 相似的功能,但 EfficientTAM 更高效。
4.3 消融研究
(1)對象指針標記的影響
我們在記憶模塊執行交叉注意力時研究了對象指針標記的影響。我們對是否包含對象指針標記的交叉注意力進行了消融實驗。我們發現,對象指針標記顯著提高了在SA-V測試數據集上的性能,從72.1 J&F提升到74.5 J&F,這與SAM 2(Ravi等人,2024年)的結果一致。這表明對象指針標記需要與記憶庫中的空間標記進行交叉注意力。
(2)記憶標記的結構
我們對記憶模塊中高效交叉注意力的記憶標記結構進行了消融實驗。在我們的高效交叉注意力中,我們利用記憶空間標記的局部性來獲得更粗略的表示,并將粗略嵌入與對象指針標記連接起來。我們觀察到,如果對整個記憶標記而不是僅對空間標記進行池化,會導致性能大幅下降,在SA-V測試集上下降了2.3 J&F。
(3)窗口大小的影響
我們在方程(5)中進行了平均池化,以獲得良好的替代品。我們嘗試了2×2和4×4的窗口大小。我們發現,將高效交叉注意力的窗口從2×2增加到4×4會導致約1 J&F的準確率下降,但速度略有提升。因此,我們使用2×2的窗口大小,在準確率和效率之間取得平衡。
(4)線性交叉注意力
我們探索了適應一種代表性高效注意力方法,例如線性注意力(Choromanski等人,2020年;Cai等人,2023年;You等人,2023年),通過利用矩陣乘法的結合律。我們發現,使用矩陣乘法的結合律的線性注意力會導致顯著的性能下降,在SA-V測試集上超過10 J&F的準確率下降,與我們提出的高效交叉注意力相比。因此,利用底層標記結構進行高效交叉注意力更為有效。
(5)高效交叉注意力變體的比較
我們比較了不同的高效交叉注意力變體。我們發現,方程(5)中的高效交叉注意力優于Linformer變體,在SA-V測試集上,74 J&F對比73.4 J&F。然而,我們發現方程(6)可以實現相當的性能,如表4所示。
(6)輸入分辨率的影響
我們對視頻對象分割的輸入分辨率進行了消融實驗。默認情況下,我們使用了1024×1024的分辨率。我們嘗試了不同的輸入分辨率,例如512×512。表5顯示,降低輸入分辨率會導致一些性能下降,但會提高效率,尤其是在移動設備上,例如在iPhone 15上速度提升了12.5倍。這為具有不同延遲和質量需求的實際部署提供了靈活性。
5 結論
我們重新審視了使用普通非層次化圖像編碼器來構建高效的視頻對象分割和追蹤任意目標模型——EfficientTAM。EfficientTAM 使用普通的輕量級 ViT 圖像編碼器,在圖像和視頻分割能力上與層次化圖像編碼器相當,同時更加高效,能夠在移動設備上部署。我們還提出了一個高效的記憶模塊,該模塊通過利用空間記憶嵌入的局部性,實現了更快的交叉注意力。高效的記憶模塊進一步優化了 EfficientTAM 在視頻分割和追蹤任意目標方面的準確率與效率的權衡。在半監督視頻對象分割、可提示視頻分割以及分割任意目標任務上的廣泛實驗一致驗證了我們的 EfficientTAM 的優勢。我們的初步工作表明,EfficientTAM 在設備上追蹤任意目標方面具有許多潛在應用。