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 在设备上追踪任意目标方面具有许多潜在应用。