人工智能

TinySAM :高效的 Segment Anything 模型

mm
TinySAM : Pushing the Boundaries for Segment Anything Model

对象分割是现代计算机视觉中一个基础且至关重要的领域,它在需要大量视觉组件的应用中发挥着至关重要的作用,例如对象定位和识别,并需要实时、快速和准确的分割。这一重要性使得对象分割成为一个持续热门的研究话题,人们在实例分割、语义分割和全景分割等领域做了大量的工作。

随着对象分割的发展,Segment Anything Model(SAM)作为一个显著的工具出现,展示了出色的分割能力,并迅速被应用于各种计算机视觉应用中。使用预训练的SAM架构的框架在下游视觉任务中取得了令人印象深刻的性能。然而,尽管SAM具有出色的分割能力和高精度,但其复杂和重量级的架构需要大量的计算能力,这阻碍了其在计算资源受限的设备上的实施。

为了解决SAM的计算挑战,研究人员开发了Tiny Segment Anything Model(TinySAM),它保留了原始框架的零样本性能,同时变得更加轻量。TinySAM使用全阶段知识蒸馏方法和在线硬提示采样来创建一个更高效的学生模型。后训练量化适应提示分割任务进一步降低了计算需求。此外,TinySAM的设计目标是分层分割,几乎将推理速度提高了一倍,而不会损害性能。

本文深入探讨了TinySAM框架,探索其基本原理、架构和与其他最先进的分割框架相比的性能。让我们更详细地探讨这些方面。

TinySAM :高效的 Segment Anything 模型

Segment Anything Model 由于其出色的分割能力和大规模的分割数据集(包含超过1100万张图片和超过10亿个图像掩码),推动了多个计算机视觉应用的快速发展。由于其在分割具有任意类别和形状的对象方面的卓越性能,SAM成为执行下游任务(如图像修复、对象跟踪、3D视觉等)的框架的基础。此外,Segment Anything Model 还提供了显著的零样本分割性能,这使得其在数据有限的行业(如医疗研究和医疗成像)中受益。

尽管SAM在下游视觉任务中的性能令人印象深刻,但其复杂的架构和高计算需求却使其难以在计算资源受限的设备上实施。为了克服这个障碍,研究人员提出了MobileSAM和FastSAM等方法,旨在开发更具计算效率的SAM模型。MobileSAM框架尝试用TinyViT架构替换重型的图像编码器,而FastSAM模型则将分割任务转换为仅有一类的实例分割任务,使用YoloV8模型。虽然这些方法在降低计算需求方面取得了一定的成功,但它们并未能在零样本下游任务中保持性能。

TinySAM或Tiny Segment Anything Model是为了在不损害零样本下游任务性能的情况下降低当前SAM模型的计算需求而提出的。TinySAM框架提出了一种全阶段知识蒸馏方法,以提高紧凑的学生网络的能力。该框架在不同阶段下将学生网络蒸馏为教师网络,并通过在线硬提示采样进一步提高性能。此外,TinySAM框架还将后训练量化适应提示分割任务,以进一步降低计算成本。

TinySAM :架构和方法

在讨论TinySAM框架的架构和方法之前,了解其前身SAM框架至关重要。SAM框架自其引入以来,在各种下游视觉和对象分割任务中展示了出色的性能、多样性和泛化能力。

SAM模型的核心由三个子网络组成:提示编码器、图像编码器和掩码解码器。提示编码器的主要目标是编码具有任意形状的掩码、输入点和框,以及自由形式的文本及其位置信息。图像编码器是一个基于视觉变换器的重型网络,用于将输入图像提取为嵌入。模型使用不同的网络处理几何和文本提示。最后,掩码解码器包含一个双向变换器,接收提示和图像编码器的输出以生成最终的掩码预测。有了数据集,SAM框架展示了出色的高质量分割能力,适用于各种形状和类别的对象。此外,SAM框架还展示了出色的零样本下游视觉任务性能,包括对象提议、边缘检测、文本到掩码预测和实例分割。由于其高质量的分割能力和灵活的提示,SAM框架成为视觉应用的基础。然而,传统SAM架构的高计算需求使得其难以在资源受限的设备上实施。

知识蒸馏

知识蒸馏是一种提高紧凑网络训练性能的重要方法。知识蒸馏方法使用教师网络的输出来监督轻量级学生网络的训练。知识蒸馏方法可以分为两类:中间特征蒸馏和网络输出蒸馏,大多数知识蒸馏研究工作都集中在图像分类任务上。

如下图所示,TinySAM框架的通用架构以及在零样本实例分割任务上的性能概述。

TinySAM框架首先实施了针对SAM框架的知识蒸馏,并使用在线硬提示采样来进一步激活蒸馏过程。然后,框架将后训练量化适应提示分割任务,并在轻量级学生网络上实施。最后,模型实施了分层分割推理模式,设计用于分割任务,几乎将推理速度提高了一倍,而不会损害性能。

全阶段知识蒸馏

如前所述,Segment Anything Model 由三个子网络组成:提示编码器、图像编码器和掩码解码器。图像编码器组件基于视觉变换器,具有高计算需求。为了解决这个问题,TinySAM框架实施了一种全阶段知识蒸馏方法,指导轻量级图像编码器从学习级别到多级知识。除了传统的损失函数外,TinySAM框架在不同阶段引入了多个蒸馏损失,如下图所示。

量化

模型量化是计算机视觉框架中的一种流行方法,用于通过将权重或激活从高带宽降低到低带宽来压缩模型,尝试在不显著降低输出质量的情况下降低计算复杂度和存储需求。

TinySAM中的量化主要是将浮点张量投影到位整数张量,使用一个缩放因子,矩阵乘法和量化矩阵之间的距离度量对于优化缩放因子至关重要。

分层分割

Segment Anything Model 提出使用自动掩码生成器,通过网格采样点来分割图像中的所有内容。然而,使用密集的点网格会导致过于细致的分割输出,并且该过程需要大量的计算资源和高昂的运营成本。此外,一方面,对象的过多采样点可能导致对象的不同部分被错误地分割为单独的掩码,而另一方面,所有内容模式的推理时间成本主要是由于图像编码器被显著压缩。为了降低所有内容模式的运营成本,TinySAM框架使用了一种分层掩码生成方法,与原始SAM框架的策略不同,如下图所示。

不同于原始SAM框架的方法,TinySAM模型仅使用每边的25%点,因此仅使用原始设置中的1/16个点。然后,模型使用这些提示推断掩码解码器和提示编码器,并获得输出。模型然后过滤掉置信度超过某个阈值的掩码,并将相应的位置标记为潜在的最终预测区域。由于模型将这些区域视为高置信度实例的分割结果,因此无需生成点提示。这种策略不仅可以防止对象的过度细分,还可以显著降低运营成本和计算需求。框架然后合并和后处理两个轮次的结果以获得最终的掩码。

TinySAM :实验和结果

为了加速蒸馏过程,TinySAM框架预先计算和存储教师网络的图像嵌入,这样就不需要在训练过程中重复计算教师网络的重型图像编码器。对于后训练量化,TinySAM框架量化了所有的矩阵乘法层、卷积层、反卷积层和线性层,模型使用通道级缩放因子进行卷积和反卷积,而使用头级缩放因子进行矩阵乘法层,并使用线性级缩放因子进行线性层。模型还在零样本下游任务上进行了评估。

在零样本实例分割任务中,TinySAM框架遵循其前身Segment Anything Model的实验设置,使用Vision Transformer Det-H或VitDet-H框架的对象检测结果进行实例分割。如下图所示,TinySAM框架在实例分割准确率和FLOPs评分方面优于现有方法。

此外,TinySAM模型在零样本实例分割任务上的定性性能如下图所示,绿色框表示框提示。

在零样本点有效掩码评估中,TinySAM模型在不同数据集上显著优于MobileSAM框架,并在使用较少点作为提示时取得了更好的结果。

如下表所示,结果总结了分层所有内容策略所实现的加速和降低计算需求的效果。模型使用相同的稳定性评分和阈值进行不同策略的公平比较,并总结如下。

最后的思考

在本文中,我们讨论了TinySAM,这是一个提出的框架,旨在推动任何分割任务的边界,并获得一个高效的模型架构,具有较低的计算需求和与原始SAM框架相当的准确率。TinySAM或Tiny Segment Anything Model保持并提供了原始框架的零样本性能。TinySAM框架首先实施了一种全阶段知识蒸馏方法,使用在线硬提示来蒸馏一个轻量级的学生模型。TinySAM框架然后将后训练量化适应提示分割任务,以进一步降低计算需求。此外,框架还旨在分层分割,几乎将推理速度提高了一倍,而不会影响性能。

专业为工程师,心为作家。 Kunal是一名技术作家,对AI和ML有着深厚的热爱和理解,致力于通过其引人入胜和信息丰富的文档来简化这些领域中的复杂概念。