TinySAM: Pushing the Envelope for Efficient Segment Anything Model
Abstract
- 近年来,任意分割模型(SAM)显示出强大的分割能力,在计算机视觉领域引起了极大的关注。以下大量工作基于预训练的SAM开发了各种应用程序,并在下游视觉任务中取得了令人印象深刻的性能。然而,SAM由沉重的架构组成,需要巨大的计算能力,这阻碍了SAM在计算受限的边缘设备上的进一步应用。为此,在本文中,我们提出了一个框架来获得(TinySAM),同时保持强大的零样本性能。我们首先提出了一种全阶段知识提取方法,该方法采用硬提示采样和硬掩码加权策略来提取轻量级学生模型。我们还将训练后量化应用于可提示的分割任务,并进一步降低计算成本。此外,提出了一种分层分割所有内容的策略,以将所有内容的推理速度提高2倍,几乎不会降低性能。有了所有这些提出的方法,我们的TinySAM导致了数量级的计算减少,并推动了有效分割任何任务的包络。在各种零样本传输任务上进行的大量实验证明了我们的TinySAM相对于对应方法的显著优势性能。经过预训练的模型和代码可在xinghaochen/TinySAM: Official PyTorch implementation of “TinySAM: Pushing the Envelope for Efficient Segment Anything Model” (github.com)。
- 论文地址:[2312.13789] TinySAM: Pushing the Envelope for Efficient Segment Anything Model (arxiv.org)
- 分割一切模型(Segment Anything Model,SAM)展现出了强大的分割万物能力,但是网络结构复杂,计算代价高,在资源受限的端侧应用中无法直接使用。同时,我们发现在SAM的Everything推理模式中,原始的网格点prompt的方式会带来巨大的计算代价。针对这个问题,我们提出了TinySAM模型,利用知识蒸馏、量化等手段对模型做了极致的压缩,同时提出了再Everything推理模式下的层次化推理策略。TinySAM在计算代价降低多个数量级的情况下,仍然保持了SAM模型强大的zero-shot分割能力。
Introduction
-
目标分割是计算机视觉领域的一项重要而基础性的任务。对象定位和验证等广泛的视觉应用依赖于准确快速的对象分割。大量的先前工作集中在分割任务上,包括语义分割、实例分割和全景分割。最近,Kirillov等人引入了一个强大的分段任意模型(SAM),以及一个包含1100万张图像上超过10亿个掩码的大规模分割数据集SA-1B。凭借强大的分割任意形状和类别物体的能力,SAM已成为许多下游任务的基础框架,如物体跟踪、图像修复和3D视觉等。此外,SAM强大的零样本分割能力使医学成像等数据较少的研究领域受益。
-
尽管SAM在下游视觉任务上取得了令人印象深刻的性能,但复杂的架构和巨大的计算成本使SAM难以部署在资源受限的设备上。在现代GPU上,1024×1024图像的SAM模型的推理时间可能长达2秒。最近的一些尝试试图获得计算效率更高的分段任意模型。例如,MobileSAM试图用TinyViT的轻量级架构取代图像编码器的重型组件。然而,它只通过在教师网络的图像嵌入监督下训练紧凑的图像编码器网络,以解耦的知识提取策略访问图像编码器网络。这种部分训练策略在没有最终掩模预测监督的情况下不可避免地导致性能衰减。FastSAM使用YOLOv8将分段任意任务转移到只有一个前景类别的实例分段任务。为了实现可提示分割的功能,FastSAM将后处理策略与实例分割网络一起应用。然而,这种重新制定的框架在下游零样本任务上无法实现与SAM相当的性能。
-
为了进一步推动高效分割任何事物模型的发展,本文提出了一个完整的框架来获得TinySAM,该框架在最大限度地保持零样本分割能力的同时,大大降低了计算成本。具体来说,我们提出了一种硬挖掘的全阶段知识提取方法来提高紧凑型学生网络的能力。在不同网络阶段的教师网络的监督下,以端到端的方式提取学生网络。提出了一种掩码加权蒸馏损失法,通过大量不同的SA-1B掩码有效地将信息从教师传递到学生。此外,提出了一种在线硬提示采样策略,使蒸馏过程更多地关注硬实例,从而提高最终性能。我们还将训练后量化应用于可提示的分割任务,并进一步降低了计算成本。此外,我们发现分割图像中的所有内容都需要巨大的计算成本,因为必须从网格提示点生成大量掩码。为此,提出了一种分层分割所有事物的策略,以在几乎没有性能下降的情况下将所有事物推理加速2倍。有了所有这些提出的方法,我们的TinySAM导致了数量级的计算减少,并推动了有效分割任何任务的包络。例如,与最初的SAM相比,TinySAM可以实现100倍的分段任何任务加速。在各种零样本传输任务上的大量实验证明了我们的TinySAM相对于同行的显著优势性能。
Related Work
Segment Anything Model
- 最近提出的分段任意模型(SAM)证明了其在对象分割和下游视觉任务中的通用性和多功能性。SAM由三个子网组成,即图像编码器、提示编码器和掩码解码器。图像编码器是一个基于重视觉 Transformer 的网络,它将输入图像提取到图像嵌入中。提示编码器设计用于对输入点、框、任意形状的掩码和具有位置信息的自由形式文本进行编码。几何提示和文本提示使用不同的网络进行处理。掩码解码器包含一个双向 transformer,它取图像编码器和提示编码器的输出来生成最终的掩码预测。与所提出的SA-1B数据集(包含1100万张高分辨率图像和超过10亿个高质量分割掩模)一起,SAM对任何类别和形状的对象都显示出令人印象深刻的高质量分割能力。此外,SAM在零样本下游视觉任务上表现出强大的通用性,包括边缘检测、目标建议、实例分割和文本到任务预测。由于其灵活的提示模式和高质量的分割能力,SAM已被视为视觉应用的基础模型。然而,SAM,尤其是图像编码器网络,由大参数组成,并且需要高计算能力来进行部署。因此,在资源有限的边缘设备上应用SAM并不容易。SAM的压缩和加速正成为一个重要的研究课题。
Knowledge Distillation
- Hinton等人提出了知识提取方法,通过教师网络的输出来监督轻量级学生网络的训练。从那时起,知识提取一直是提高紧凑网络在训练过程中性能的重要方法。知识提取方法大致可分为两类,即网络输出的提取和中间特征的提取。大多数知识提取方法的研究都集中在图像分类任务上。随后的工作[Learning efficient object detection models with knowledge distillation,Optical flow distillation: Towards efficient and stable video style transfer,Towards efficient instance segmentation with hierarchical distillation,Distilling object detectors via decoupled features,Structured knowledge distillation for semantic segmentation]提出了用于高级计算机视觉任务的知识提取方法,如对象检测和语义分割。张等人提出使用蒸馏方法来获得有效的分段任意模型(MobileSAM)。然而,MobileSAM仅在原始SAM的图像和相应图像嵌入的监督下访问图像编码器网络。这种部分提取策略可能会导致相当大的性能下降,因为无论是教师网络还是标记数据,都没有针对轻量级学生网络的掩码级信息的指导。
Quantization
- 模型量化也是常用的模型压缩方法之一,它将权重或激活从较高的位宽量化到较低的位宽,以在有限的精度下降的情况下降低存储要求和计算复杂性。有两种类型的模型量化方法,量化感知训练(QAT)和训练后量化(PTQ)。QAT方法需要标记的训练数据集和大量的训练成本,而PTQ方法只需要一个小的未标记的校准数据集,因此更高效。已经提出了许多先前的PTQ方法来搜索卷积神经网络的适当量化参数。由于视觉 transformer 在各种视觉任务上取得了显著的性能,最近的工作[Q-vit,Post-training quantization for vision transformer,Correlation congruence for knowledge distillation,Ptq4vit]研究了如何将训练后量化应用于基于ViT的模型,并在8位量化配置下取得了强大的性能。然而,对于可提示的分割任务的量化,尤其是对于任何分割模型,仍然没有探索。
Methodology
Overview of TinySAM
- 本文提出了一个获得高效SAM的框架,如下图所示。首先,我们在第3.2节中介绍了专门为SAM设计的硬挖掘全阶段知识蒸馏。为了进一步激活提取过程,所提出的硬掩模加权和硬提示采样策略都被用于挖掘从教师网络到学生网络的基本知识。其次,第3.3节中描述的后训练量化方法适用于可提示的分割任务,并应用于轻量级学生网络。第三,设计了3.4中描述的分层 everything 推理模式来分割任何任务,该模式可以避免大量冗余计算,而精度损失可以忽略不计,并将推理时间加快2倍。
-
-
(a)我们提出的方法的总体框架。由硬挖掘全阶段知识提取、训练后量化和分层推理模块组成,计算成本按数量级降低。(b) 所提出的TinySAM可以在保持性能的同时节省可观的计算成本。在NVIDIA T4 GPU上使用TensorRT测试延迟。
-
Hard Mining Full-Stage Knowledge Distillation
-
SAM由三个子网组成,即图像编码器、提示编码器和掩码解码器。图像编码器网络是基于视觉 Transformer 的,并且消耗很大的计算成本。受MobileSAM的启发,我们使用轻量级的TinyViT来取代原来的重型图像编码器网络。这种简单的替代存在相当大的性能衰减。因此,我们提出了一种硬挖掘全阶段知识提取策略,从多个知识层次指导轻量级图像编码器的学习过程。除了预测结果和地面实况标签之间的传统损失外,我们还介绍了不同阶段的多重蒸馏损失,如下图所示。
-
-
硬挖掘的框架全阶段知识提炼。对于SA-1B数据集的海量掩码,我们设计了提示的硬提示采样和蒸馏损失的硬掩码加权。对于硬提示采样过程,星形表示具有不同迭代的采样点。随着提示采样的迭代,采样区域更靠近目标掩码的边缘,这使得学生网络相对更难学习提示。此外,根据学生和教师网络之间的差距,在计算蒸馏损失时,为每个mask分配不同的权重。
-
-
具体来说,我们从多个知识层面选择教师网络的几个节点来指导学生网络的学习。首先,我们选择图像编码器的输出特征,即图像嵌入,作为提取信息。图像嵌入集中了输入图像中的信息,这是预测过程中的基础知识。对于 I 的输入图像,用于图像嵌入的蒸馏损失函数可以表示为,
-
L e m b e d d i n g = L ( E i m g T ( I ) , E i m g S ( I ) ) L_{embedding}=L(E^T_{img}(I),E^S_{img}(I)) Lembedding=L(EimgT(I),EimgS(I))
-
其中 E i m g S E^S_{img} EimgS和 E i m g T E^T_{img} EimgT分别表示用于学生和教师网络的图像编码器。由于图像级信息与掩模预测不直接相关,因此更接近最终输出的特征对于该分割任务至关重要。自然地,教师网络的最终输出被选择为一个蒸馏点。输出蒸馏损失 L o u t p u t L_{output} Loutput可以描述为,
-
L o u t p u t = L ( D m a s k T ( E i m g T ( I ) , q ) , D m a s k S ( E i m g S ( I ) , q ) ) L_{output} = L(D^T_{mask}(E^T_{img}(I), q), D^S_{mask}(E^S_{img}(I), q)) Loutput=L(DmaskT(EimgT(I),q),DmaskS(EimgS(I),q))
-
其中 D m a s k S D^S_{mask} DmaskS 和 D m a s k T D^T_{mask} DmaskT分别是用于学生和教师的掩码解码器。q 表示掩码解码器的查询,它是提示嵌入和输出 token 的级联。由于SAM的结构相当复杂,前面提到的两种蒸馏可能不一致,对于轻量级学生来说很难学习。我们进一步提出从掩码解码器的双向 Transformer 中提取输出标记,该 Transformer 与来自提示嵌入和图像嵌入的信息相互作用。它以更抽象的方式捕获目标掩码信息。相应的蒸馏损失 L t o k e n L_{token} Ltoken 可以描述为,
-
L t o k e n = L ( T T ( E i m g T ( I ) , q ) , T S ( E i m g S ( I ) , q ) ) L_{token} = L(T^T (E^T_{img}(I), q), T^S (E^S_{img}(I), q)) Ltoken=L(TT(EimgT(I),q),TS(EimgS(I),q))
-
其中 T S T^S TS和 T T T^T TT是掩码解码器的双向 Transformer 模块,L表示损失函数。我们经验发现,特征差的数值可能会使传统使用的MSE损失(ℓ2距离)太小而不能很好地优化。因此,我们使用ℓ1距离函数。总的蒸馏损失函数 L d i s t r a c t L_{distract} Ldistract可以表示为,
-
L d i s t i l l = α ∗ L e m b e d d i n g + β ∗ L t o k e n + γ ∗ L o u t p u t L_{distill} = α ∗ L_{embedding} + β ∗ L_{token} + γ ∗ L_{output} Ldistill=α∗Lembedding+β∗Ltoken+γ∗Loutput
-
其中,α、β、γ表示每个蒸馏损失的超参数。总训练损失是掩模预测 L m a s k L_{mask} Lmask和IoU预测 L i o u s L_{ious} Lious的蒸馏损失和 GT 损失的线性组合,其中 L m a s k L_{mask} Lmask是焦点损失和dice损失的组合, L i o u s L_{ious} Lious是预测的IoU和计算的IoU之间的L1损失函数。
-
L t o t a l = L d i s t i l l + L m a s k + L i o u s . L_{total} = L_{distill} + L_{mask} + L_{ious}. Ltotal=Ldistill+Lmask+Lious.
-
-
Hard Mask Weighting. 为了使知识提取更加有效,我们在计算损失时设计了一种硬掩模加权策略。有一个观察结果表明,由于细粒度和没有语义约束,SA-1B 数据集的单个图像中的掩码可能非常不同。如上图所示,分割具有复杂边界的标志可能很困难,而分割具有高对比度颜色的矩形窗口可能很容易。硬 mask 应合理地分配更大的重量,供学生学习。具体来说,我们计算学生和教师网络输出的差距,以指示掩模硬度 H i H_i Hi。
-
H i = s i g m o i d ( I o U ( M i T , M i G T ) I o U ( M i S , M i G T ) + ϵ − 1 ) H_i=sigmoid(\frac{IoU(M^T_i,M^{GT}_i)}{IoU(M^S_i,M^{GT}_i)+ ϵ }-1) Hi=sigmoid(IoU(MiS,MiGT)+ϵIoU(MiT,MiGT)−1)
-
其中 M i T 、 M i S 、 M i G T M^T_i、M^S_i、M^{GT}_i MiT、MiS、MiGT分别表示学生网络的掩码预测、教师网络的掩码预报和第 i 个掩码的 GT。因此,蒸馏损失可以更新为
-
L d i s t i l l ∗ = α ∗ L e m b e d d i n g + β ∗ L t o k e n + γ ∗ ∑ i = 1 N H i ∗ L o u t p u t i L ^∗_{distill} = α∗L_{embedding}+β∗L_{token}+γ∗ \sum_{i=1}^N H_i∗L^i_{output} Ldistill∗=α∗Lembedding+β∗Ltoken+γ∗i=1∑NHi∗Loutputi
-
-
Hard Prompt Sampling. 通常,可以从标记的训练数据中随机采样来生成提示,以驱动基于提示的掩码预测网络SAM的端到端训练。为了进一步简化教师和轻量级学生网络之间的蒸馏学习过程,我们提出了一种硬提示采样策略,使训练样本集中在预测的困难区域。以点提示为例,最初在标记的掩模区域 M g t M_{gt} Mgt 内对点 P 0 P_0 P0 进行采样。将这些初始点与输入图像一起输入到网络中,以获得预测掩模区域 M 0 M_0 M0。然后,我们从 M g t M_{gt} Mgt和 M 0 M_0 M0的差集中采样提示点,并交互执行该过程。从 M g t M_{gt} Mgt和 M i M_i Mi的差集对第(i+1)轮采样点Pi进行采样,即:
-
P i + 1 ∈ M g t − M i , i = 0 , 1 , 2 , . . . M i = D m a s k ( E p r o m p t ( P i ) , E i m g ( I ) ) . P_{i+1} ∈ M_{gt} − M_i, i = 0, 1, 2, ...\\ M_i = D_{mask}(E_{prompt}(P_i),E_{img}(I)). Pi+1∈Mgt−Mi,i=0,1,2,...Mi=Dmask(Eprompt(Pi),Eimg(I)).
-
当应用于训练过程时,第 i 次迭代是从0到9随机采样的,这使得在约束范围内采样提示变得困难。上图的底部显示了采样提示的位置随迭代的变化,绿色星号表示具有在线硬提示采样策略的新采样点提示。迭代次数越多,采样点就越接近 GT 遮罩的边缘区域。
-
Quantization
-
量子化旨在将浮点张量 x 投影到具有缩放因子s 的bbit整数张量 x q x_q xq。均匀对称量子化可以公式化如下,
- x q = Q ( b , s ) = c l i p ( r o u n d ( x s ) , − 2 b − 1 , 2 b − 1 − 1 ) . x_q = Q(b, s) = clip(round( \frac x s ), −2 ^{b−1} , 2^{b−1} − 1). xq=Q(b,s)=clip(round(sx),−2b−1,2b−1−1).
-
对于矩阵乘法O=AB,它可以用两个缩放因子 s A s_A sA和 s B s_B sB来量化,并且量化的矩阵表示为 O ˉ = A ˉ B ˉ \bar O=\bar A\bar B Oˉ=AˉBˉ。测量 O ˉ \bar O Oˉ和O之间距离的度量对于优化 A q A_q Aq和 B q B_q Bq至关重要。继量化方法在图像分类模型中的成功实践之后,我们将 hessian 引导度量作为距离来求解缩放因子,这与任务损失更为一致。与分类任务不同,SAM的可提示分割任务输出包含细粒度掩码的分割预测。因此,我们使用掩码和IOU的Kullback-Leible(KL)散度作为任务损失,并使用一些校准数据来计算hessian 矩阵,任务损失公式为:
-
L = K L ( ˆ y p r e d , y p r e d ) + K L ( ˆ y i o u , y i o u ) , L = KL(ˆy_{pred}, y_{pred}) + KL(ˆy_{iou}, y_{iou}), L=KL(ˆypred,ypred)+KL(ˆyiou,yiou),
-
其中, y p r e d y_{pred} ypred和 y i o u y_{iou} yiou是浮点模型的输出, y ˉ p r e d \bar y_{pred} yˉpred和 y ˉ i o u \bar y_{iou} yˉiou是量化后的输出。
-
-
在指定距离度量后,我们可以将 s A s_A sA和 s B s_B sB作为交替迭代网格搜索问题来求解。使用校准数据,我们获得A和B的最大值,分别为 A m a x A_{max} Amax和 B m a x B_{max} Bmax,并使用两个参数α和β来指定 s A s_A sA和 s B s_B sB的搜索范围, [ α A m a x 2 b − 1 , β A m a x 2 b − 1 ] [α\frac{A_{max}} {2^{b−1}},β\frac{A_{max}} {2^{b−1}}] [α2b−1Amax,β2b−1Amax]和 [ α B m a x 2 b − 1 , β B m a x 2 b − 1 ] [α\frac{B_{max}} {2^{b−1}},β\frac{B_{max}} {2^{b−1}}] [α2b−1Bmax,β2b−1Bmax]。这两个搜索范围被分别线性地划分为n个候选选项。 A q A_q Aq和 B q B_q Bq被交替地优化若干轮。
-
softmax之后的矩阵乘法输入在区间[0,1]的两端不均匀分布,而GELU之后的特征在正负范围之间变化很大。这两种情况与均匀量化的假设相去甚远,即神经网络中的激活服从高斯分布。违反将导致高量化误差。因此,我们将特征分为两组,并使用两个缩放因子来减少量化误差。
Hierarchical Segmenting Everything
-
SAM提出了一种自动掩码生成器,它将点采样为网格来分割所有内容。然而,我们发现密集点网格导致分割结果过于细粒度,也占用了大量的计算资源。一方面,对于一个完整的对象,过多的采样点可能会导致对象的稍微不同的部分被错误地分割为单独的遮罩。另一方面,由于图像编码器已经被很大程度地缩小,所以所有模式推断的时间成本主要在掩码解码器部分。对于SAM自动掩码生成器的默认设置,它采样32×32=1024个点作为提示,这意味着掩码解码器被推断1024次。在单个V100 GPU上,图像编码器的成本为16ms,掩模解码器的成本为894ms。
-
为了降低 everything 模式的时间成本,我们提出了一种分层掩码生成方法。我们的分层策略与原始策略之间的比较如下图所示。与原始的一切模式不同,在第一步中,我们只在每侧使用25%的点,因此总点是原始设置的1/16,如图(c)所示。然后,我们根据这些提示推断出提示编码器和掩码解码器,并得到结果。
-
-
我们的分级策略与原始策略之间的比较。(a) 原始everything 模式的点采样(以每侧16点为例)。(b) 原始策略的分割结果。(c) 我们分级策略的第一步,仅对1/16个点进行采样。(d) 从(c)中获得高置信度区域,忽略该区域中的点。高置信度区域显示为白色遮罩。(e) 我们分层策略的细分结果。
-
-
然后,我们过滤掉一些置信度超过阈值 τ 的掩码,并将相应的位置标记为可以被视为最终预测的区域。对于这些区域,由于它们被认为是具有高置信度的实例的分割结果,因此不需要重新生成点提示。然后,我们使用原始设置对密度相同的点进行采样,但忽略上面区域中的点。如图(d)所示,我们的策略忽略了第一头牛的草和身体上的大多数点。同时,保留第二头牛和天空上的点以进一步分割。具体地说,第二头奶牛在第一轮中被错误地分割为与第一头奶牛相同的对象。这种策略可以避免推理时间的成本和对象的过细粒度分割。然后,我们利用在第二轮中采样的点提示来获得掩码预测。最后,将这两轮的结果合并并进行后处理,得到最终的掩码。我们的方法忽略了50%以上的点,从而大大减少了延迟。
Experiments
Implementation Details
- 我们使用TinyViT-5M作为轻量级学生图像编码器,使用SAM-H作为教师模型,遵循先前的工作[mobilesam]。部分SA-1B数据集被用作全阶段蒸馏的训练数据。我们采用了Adam优化器,对学生网络进行了8个epoch的训练。对于每次迭代,我们根据第3.2节中描述的硬提示采样策略对64个提示进行采样。为了加速提取过程,预先计算并存储了来自教师网络的图像嵌入。因此,教师网络的重图像编码器不需要在训练时间内重复计算。对于训练后量化,我们设置α=0.01,β=1.2,n=100,轮次=3用于迭代搜索。我们使用8幅图像在SA-1B数据集上校准量化模型。我们对实例分割和点提示分割等下游任务进行零样本评估。根据SAM的建议,采用多输出模式,并且最终的掩码预测是具有最高IoU预测的掩码预测。
Zero-Shot Instance Segmentation
-
对于零样本实例分割任务,我们严格遵循SAM的实验设置,并使用ViTDet-H的对象检测结果作为实例分割的框提示。我们在COCO数据集和LVIS v1的基准上评估了模型的零样本实例分割任务。我们将TinySAM与SAM的不同变体以及之前的高效模型(如FastSAM、MobileSAM和EfficientSAM)进行了比较。如下表所示,与现有方法相比,所提出的TinySAM获得了优越的性能。具体而言,我们的TinySAM在FLOP和实例分割精度方面优于FastSAM,即,仅12.2%的FLOP和25%的延迟,AP提高了约4%。
-
-
在COCO和LVIS v1数据集上的零样本实例分割结果。Zero-shot transfer方法是提示来自全监督ViTDet模型的检测框。TinySAM和量化Q-TinySAM在平均精度上表现出有利的性能。延迟在NVIDIA T4 GPU上进行了测试。
-
-
在相同的计算成本下,我们的TinySAM在COCO数据集上也分别比MobileSAM和LVIS v1数据集上实现了1.3%以上的AP和1.9%以上的AP。在COCO数据集上具有类似性能的情况下,TinySAM比EfficientSAM快2倍。我们的TinySAM的W8A8量化变体(Q-TinySAM)也在不同的方法中获得了有竞争力的性能。具体而言,与MobileSAM相比,QTinySAM在COCO数据集上实现了0.4%以上的AP,仅具有48%的FLOP和62.5%的延迟。COCO验证集上的一些可视化结果如下图所示。与SAM的其他高效变体相比,我们提出的TinySAM捕获了更清晰、更平滑的边界。
-
-
零样本实例分割的COCO验证数据集的可视化结果。绿色框标记框提示。我们的TinySAM捕捉到更清晰、更平滑的边界,特别是对于小尺寸或类似纹理特征的硬目标。
-
Zero-shot Points Valid Mask Evaluation
-
在本节中,我们还评估了TinySAM在根据提示从多个点分割对象方面的性能。我们使用与先前工作[Lvis,Segment anything]相同的点选择度量,该度量计算假阳性和假阴性掩码的距离变换,然后以最大值对点进行采样。我们计算每个数据集的mIoU,以评估不同模型的性能。
-
我们从[sam]中使用的总共23个数据集中选择一个子集进行有效评估,其中包括BBBC038v1、DOORS、TimberSeg和LVIS。为了进行公平的比较,我们按照SAM论文中下表的设置对图像和掩模进行采样,并在评估中使用相应分割中的前N个掩模。
-
评估结果如下图所示。我们的TinySAM在LVIS和TimberSeg数据集上显著优于MobileSAM,并在DOORS数据集上获得类似的性能。此外,当使用较少的点数作为提示时,我们的TinySAM在BBBC038v1上获得了更好的结果。我们还报告了所有四个数据集的平均IoU,如图中的右图所示。所提出的TinySAM优于MobileSAM,并且获得了与SAM-B接近的性能。
-
-
零样本点有效掩模的评价结果。X轴表示提示中使用的点数,Y轴表示评估的所有遮罩的mIoU。所提出的TinySAM优于MobileSAM,并实现了接近SAM ViT-B的结果。
-
Everything Mode Acceleration
- 我们在COCO验证集上评估了我们提出的分层一切推理策略。延迟基准测试在单一NVIDIA V100 GPU上进行,适用于所有模式。我们从val 2017中采样了100张img id最小的图像,并对这些样本进行了万物模式推断。所有模式中使用的阈值都与默认值相同。结果如下表所示。我们将相同的阈值和稳定性分数应用于用不同策略评估的同一模型,以进行公平的比较。但这些模型之间可能有所不同。与原来的32×32点网格策略相比,我们的分层策略获得了可比的结果,同时推理时间成本降低了约50%。
-
-
原始点网格策略和我们的分层策略的比较。对COCO val2017套装前100张图像的评估。
-
- 下图显示了点网格策略和所提出的分层策略之间的视觉比较。我们可以发现,在某些情况下,所提出的分层策略避免了对象的过度分割,这弥补了IoU的一些损失。高效的everything模式推断的更直观的结果如下图所示。
-
-
COCO val2017集合上所有模式分割结果的可视化比较。(a) 点网格策略结果,mIoU 0.5292。(b) 分层策略结果,mIoU 0.6429。在这种情况下,我们的方法避免了对对象的过度细化分割。
-
Ablation Studies
-
在本节中,我们对所提出的方法在COCO验证数据集上的零样本实例分割任务进行了消融研究。实验设置与第4.2节所述相同。
-
不同模块的影响。我们首先评估了不同模块的影响,即全阶段知识蒸馏损失、在线硬提示采样和后量化。如下表所示,利用我们提出的全阶段蒸馏策略将性能从40.7%提高到41.4%。结合在线硬提示采样策略,我们的方法可以获得0.5%的AP增益。在硬掩模加权损失的情况下,性能可以进一步提高到42.3%。使用训练后量化导致0.9%的AP退化,但也大大降低了计算成本,这也在主要实验中进行了讨论。
-
-
分别评估了蒸馏损失、在线硬提示采样和量化对COCO验证数据集零样本实例分割的影响。
-
-
不同蒸馏损失的影响。对于详细的全阶段知识蒸馏过程,我们从教师网络中研究了所提出的三级蒸馏的必要性。下表显示了不同蒸馏损失组合的消融结果。输出蒸馏损失由于接近监督信息,并且与教师网络的相似性直接反映在评估度量中,因此它起着重要作用。token loss和嵌入 loss 都是有益的,因为它们与教师网络的关键节点有关,分别反映了图像级别的信息和提示与图像的交互。输出损耗的硬掩模加权可以进一步提高性能。
-
-
COCO 测试集零样本实例分割的知识蒸馏损失组合消融研究。
-
-
分层策略的点密度和阈值。在下表中,我们在点密度和高置信度掩模阈值 τ 的不同设置下进行消融研究。更多的点和更高的阈值 τ 导致更精确的结果,但推理时间更长。与第一轮相比,第二轮的点密度更敏感。考虑到准确性和效率,粗体的设置是一个很好的平衡,并用于 everything 推理模式的其他实验。
-
-
分层策略的点密度和阈值消融。
-
-
用于量化的不同比特。我们在这里探讨不同位宽的影响。下表报告了COCO数据集的平均精度。从结果中,我们可以得出结论,量化到8比特只会导致轻微的性能下降。我们还通过将量化比特宽度进一步减小到6比特来展示性能。
-
-
COCO验证数据集上零样本实例分割不同量化位宽的消融研究。
-
Conclusion
- 在本文中,我们提出了一个框架来突破分段任意任务的极限,并获得了一个高效的模型TinySAM。我们首先提出了一种全阶段知识提取方法,该方法采用硬掩模加权和硬提示采样策略来提取轻量级学生模型。我们还将训练后量化应用于可提示的分割任务,并进一步降低了计算成本。此外,提出了一种分层分割所有内容的策略,以将所有内容的推理速度提高2倍,几乎不会降低性能。有了所有这些提出的方法,我们的TinySAM导致了数量级的计算减少,并推动了高效分段任何任务的包络。在各种零样本传输任务上进行的大量实验证明了我们的TinySAM相对于对应方法的显著优势性能。我们希望所提出的TinySAM为设计高效的分段任意模型带来有益的前景。
Appendix
-
我们为补充材料提供了更多的可视化结果。下图显示了LVIS v1数据集上的零样本实例分割。为了清晰呈现,只有VitDet-H检测到的得分高于0.8的方框才会被提示到模型中。LVIS数据集比COCO数据集具有更多的细粒度标签,在这一点上,考虑到准确性和效率,所提出的TinySAM表现出更大的优势。
-
-
零样本实例分割的LVIS v1数据集的可视化结果。绿色框表示框提示。与FastSAM和MobileSAM相比,TinySAM生成更准确、更平滑的掩码。
-
-
下图显示了所提出的具有分层 everything 推理及其对应算法的TinySAM模型的万物推理结果。TinySAM捕捉清晰的边界并产生更细粒度的掩模,而MobileSAM和FastSAM有时会产生制造的边界和掩模。TinySAM显示出与原始SAM更接近的性能,同时消耗的计算成本显著降低。
-
-
TinySAM模型与分层 everything 推理及其对应算法的可视化结果。与FastSAM和MobileSAM相比,TinySAM捕获细粒度边界和掩码,与计算成本高昂的SAM-H模型表现出相似的性能。
-