0. 论文信息
标题:TinySAM: Pushing the Envelope for Efficient Segment Anything Model
作者:TinySAM: Pushing the Envelope for Efficient Segment Anything Model
机构:University of Science and Technology of China、Huawei Noah’s Ark Lab
原文链接:https://arxiv.org/abs/2312.13789
代码链接:https://github.com/xinghaochen/TinySAM
1. 导读
最近,分割任何东西模型(SAM)显示出强大分割能力,在计算机视觉领域引起了极大的关注。大量后续工作基于预训练的SAM开发了各种应用,并在下游视觉任务中取得了令人印象深刻的性能。然而,SAM结构庞大,计算量大,阻碍了其在计算受限的边缘设备上的进一步应用。为此,在本文中,我们提出了一个框架,以获得微小片段任何模型(TinySAM ),同时保持强大的零镜头性能。我们首先提出一种全阶段知识提取方法,使用硬提示抽样和硬掩码加权策略提取轻量级学生模型。我们还使训练后量化适应基于提示的分割任务,并进一步降低计算成本。此外,还提出了一种分层分割万事万物的策略,通过2×几乎没有性能下降。通过所有这些建议的方法,我们的TinySAM导致了数量级的计算减少,并推动了高效分段任何任务的包络。在各种零射击转移任务上的大量实验证明了我们的TinySAM相对于对应方法的显著优势性能。
2. 效果展示
(a)我们提出的方法的总体框架。包括全阶段知识蒸馏、后训练量化和分层一切推理的模块计算成本大大减少。(b)提出的TinvSAM在保持性能的同时可以节省可观的计算成本。延迟在NVIDIAT4 GPU上使用TensorRT进行测试。
过程层次化一切策略的可视化。图(a)显示了高置信度区域的第一批稀疏提示点后的中间结果带有白色掩模,并保留了第二批密集的提示点,带有绿色星星。图(b)显示了最终分割结果,小物体可以准确分割。
零样本实例分割的COCO验证数据集(上3行)和LVIS1数据集(下3行)的可视化结果。绿色框标记了ViTDet-H检测器的框提示。TinySAM能够捕捉到更清晰、更平滑的边界,特别是对于小尺寸或类似纹理特征的困难目标。
3. 引言
目标分割是计算机视觉领域一项重要且基础的任务。诸如目标定位和验证等广泛的视觉应用均依赖于准确且快速的目标分割。大量先前的工作聚焦于分割任务,包括语义分割、实例分割和全景分割。最近,Kirillov等人引入了一个强大的“分割任意事物”模型(SAM),以及一个包含超过1100万张图像上10亿多个掩码的庞大分割数据集SA-1B。SAM具有分割任意形状和类别的目标的强大能力,已成为许多下游任务(如目标跟踪、图像修复和3D视觉等)的基础框架。此外,SAM强大的零样本分割能力为医疗成像等数据量较少的研究领域带来了益处。尽管SAM在下游视觉任务上取得了令人瞩目的性能,但其复杂的架构和巨大的计算成本使得SAM难以在资源受限的设备上部署。SAM模型对1024×1024图像的推理时间在现代GPU上可能长达2秒。最近的一些尝试试图获得一个计算效率更高的“分割任意事物”模型。例如,MobileSAM试图用TinyViT的轻量级架构替换图像编码器的繁重组件。然而,它仅通过用教师网络的图像嵌入进行监督来训练紧凑的图像编码器网络,从而采用了一种解耦的知识蒸馏策略来访问图像编码器网络。这种部分训练策略在没有最终掩码预测的监督下,不可避免地会导致性能下降。FastSAM将“分割任意事物”任务转换为只有一个前景类别的实例分割任务,使用的是YOLOv8。为了实现基于提示的分割功能,FastSAM将后处理策略与实例分割网络相结合。然而,这个重新构建的框架在下游零样本任务上的性能无法与SAM相媲美。
推荐课程:
彻底搞懂大模型数学基础剖析、原理与代码讲解
。
为了进一步推动高效“分割任意事物”模型的发展,本文提出了一个完整框架来获得TinySAM,该框架在最大限度地保持零样本分割能力的同时,大大降低了计算成本。具体而言,我们提出了一种硬挖掘全阶段知识蒸馏方法,以提高紧凑型学生网络的性能。学生网络在端到端的方式下,在不同网络阶段的教师网络的监督下进行蒸馏。我们提出了一种掩码加权蒸馏损失,以有效地通过大量不同的SA-1B掩码将信息从教师传递给学生。此外,还提出了一种在线硬提示采样策略,以使学生网络具备分割任意形状和类别的目标的强大能力,SAM已成为许多下游任务(如目标跟踪、图像修复和3D视觉等)的基础框架。此外,SAM强大的零样本分割能力为医疗成像等数据量较少的研究领域带来了益处。尽管SAM在下游视觉任务上取得了令人瞩目的性能,但其复杂的架构和巨大的计算成本使得SAM难以在资源受限的设备上部署。SAM模型对1024×1024图像的推理时间在现代GPU上可能长达2秒。最近的一些尝试试图获得一个计算效率更高的“分割任意事物”模型。例如,MobileSAM试图用TinyViT的轻量级架构替换图像编码器的繁重组件。然而,它仅通过用教师网络的图像嵌入进行监督来训练紧凑的图像编码器网络,从而采用了一种解耦的知识蒸馏策略来访问图像编码器网络。这种部分训练策略在没有最终掩码预测的监督下,不可避免地会导致性能下降。FastSAM将“分割任意事物”任务转换为只有一个前景类别的实例分割任务,使用的是YOLOv8。为了实现基于提示的分割功能,FastSAM将后处理策略与实例分割网络相结合。然而,这个重新构建的框架在下游零样本任务上的性能无法与SAM相媲美。
为了进一步推动高效“分割任意事物”模型的发展,本文提出了一个完整框架来获得TinySAM,该框架在最大限度地保持零样本分割能力的同时,大大降低了计算成本。具体而言,我们提出了一种硬挖掘全阶段知识蒸馏方法,以提高紧凑型学生网络的性能。学生网络在端到端的方式下,在不同网络阶段的教师网络的监督下进行蒸馏。我们提出了一种掩码加权蒸馏损失,以有效地通过大量不同的SA-1B掩码将信息从教师传递给学生。此外,还提出了一种在线硬提示采样策略,以使学生蒸馏出
4. 主要贡献
本文提出了一个框架以获得高度高效的SAM。首先,我们为SAM引入了一种专门设计的硬挖掘全阶段知识蒸馏方法。为了进一步激活蒸馏过程,我们采用了提出的硬掩码加权和硬提示采样策略,以从教师网络向学生网络挖掘关键知识。其次,我们针对基于提示的分割任务,对后训练量化方法进行了调整,并将其应用于轻量级学生网络。第三,我们设计了一种分层全事物推理模式,用于全事物分割任务,该模式可以在几乎不损失精度的情况下避免大量冗余计算,并将推理时间加快2倍。
5. 方法
全阶段知识蒸馏框架。对于SA-1B数据集的大量掩码,我们为提示设计硬提示采样,为蒸馏损失设计硬掩码加权。在采样过程中,星号表示具有不同迭代次数的采样点。随着迭代次数的增加,采样区域更接近目标掩码的边缘,这交使得提示对学生网络的学习相对更难。此外,计算蒸馏损失时,根据学生网络和教师网络之间的差距,为每个掩码分配不同的权重。
我们的分层策略与原始策略的对比。(a)原始所有模式中的点采样(以每边16个点为例)()原始策略的分段结果(c)我们分层策略的第一步,仅采样1/16个点。(d)从(c)中获取高置信度区域并忽略该区域内的点。高置信度区域显示为白色掩模。(e)我们分层策略的分段结果。
6. 实验结果
7. 总结 & 未来工作
在这篇论文中,我们提出了一种框架来推动名为TinySAM的先进模型。我们首先提出了一种全阶段知识蒸馏方法,使用硬掩码权重和硬提示采样策略来蒸馏一个轻量级学生模型。我们还适应了后训练量化到基于提示的分割任务,进一步降低了计算成本。此外,我们提出了一种分层分割一切的策略,通过几乎不降低性能的方式将一切推理速度提高2倍。通过所有这些拟议的方法,我们的TinySAM实现了数量级的计算减少,并推动了高效分割任何任务的前沿。在各种零样本迁移任务上的广泛实验表明,我们的TinySAM相对于对应方法具有显著的优势性能。我们希望提出的TinySAM为设计高效分割任何模型带来有益的视角。
对更多实验结果和文章细节感兴趣的读者,可以阅读一下论文原文~
本文仅做学术分享,如有侵权,请联系删文。
3D视觉交流群,成立啦!
目前我们已经建立了3D视觉方向多个社群,包括
2D计算机视觉
、
最前沿
、
工业3D视觉
、
SLAM
、
自动驾驶
、
三维重建
、
无人机
等方向,细分群包括:
工业3D视觉
:相机标定、立体匹配、三维点云、结构光、机械臂抓取、缺陷检测、6D位姿估计、相位偏折术、Halcon、摄影测量、阵列相机、光度立体视觉等。
SLAM
:视觉SLAM、激光SLAM、语义SLAM、滤波算法、多传感器融合、多传感器标定、动态SLAM、MOT SLAM、NeRF SLAM、机器人导航等。
自动驾驶
:深度估计、Transformer、毫米波|激光雷达|视觉摄像头传感器、多传感器标定、多传感器融合、3D目标检测、路径规划、轨迹预测、3D点云分割、模型部署、车道线检测、Occupancy、目标跟踪等。
三维重建
:3DGS、NeRF、多视图几何、OpenMVS、MVSNet、colmap、纹理贴图等
无人机
:四旋翼建模、无人机飞控等
2D计算机视觉
:图像分类/分割、目标/检测、医学影像、GAN、OCR、2D缺陷检测、遥感测绘、超分辨率、人脸检测、行为识别、模型量化剪枝、迁移学习、人体姿态估计等
最前沿
:具身智能、大模型、Mamba、扩散模型、图像/视频生成等
除了这些,还有
求职
、
硬件选型
、
视觉产品落地、产品、行业新闻
等交流群
添加小助理: cv3d001,备注:
研究方向+学校/公司+昵称
(如
3D点云+清华+小草莓
), 拉你入群。
▲长按扫码添加助理:cv3d001
3D视觉工坊知识星球