华为诺亚实验室+中科大提出TinySAM | 比SAM小10倍，但是依旧完成了精度的超车！

最近，Segment Anything Model (SAM) 已经展示出了强大的分割能力，在计算机视觉领域引起了广泛关注。基于预训练的 SAM 的大量研究工作已经开发了各种应用，并在下游视觉任务上取得了令人印象深刻的性能。然而，SAM 包含重的架构，需要大量的计算能力，这阻碍了 SAM 在计算受限的边缘设备上的进一步应用。

为此，在本文中提出了一种框架，在保持强大的零样本性能的同时，获得一个微小的分割 anything 模型（TinySAM）。作者首先提出了一种全阶段知识蒸馏方法，采用在线硬提示采样策略来蒸馏一个轻量级的的学生模型。作者还适应了后训练量化到可提示的分割任务，并进一步降低了计算成本。

此外，作者提出了一种分层的分割 everything 策略，以将 everything 推理加速 2 倍，几乎不降低性能。通过所提出的所有方法，作者的 TinySAM 导致计算减少了数个数量级，并推动了高效分割 anything 任务的极限。在各种零样本迁移任务上的广泛实验表明，作者的 TinySAM 与对应方法相比具有显著的优势。

预训练模型和代码:

Pytorch：https://github.com/xinghaochen/TinySAMMindS
pore：https://gitee.com/mindspore/models/

1 Introduction

物体分割是计算机视觉领域的一个重要且基础的任务。广泛的视觉应用，如物体定位和验证，都依赖于准确的快速物体分割。许多先前的研究工作都关注于分割任务，包括语义分割，实例分割和全景分割。最近，Kirillov 引入了一个强大的分割 anything 模型（SAM），以及一个大规模的分割数据集 SA-1B，该数据集包含了 11 亿张图像上的超过 10 亿个Mask。具有任意形状和类别的目标的强大的分割能力，SAM 已成为许多下游任务的基础框架，如物体跟踪，图像修复 和 3D 视觉。此外，SAM 的强大零样本分割能力使医学影像等数据较少的研究领域受益。

尽管 SAM 在下游视觉任务上取得了令人印象深刻的性能，但复杂的架构和巨大的计算成本使得 SAM 在资源受限的设备上部署起来困难。SAM 模型对 1024×1024 图像的推理时间在现代 GPU 上可以达到 2 秒。一些最近的努力试图获得更高效的分割 anything 模型。例如，MobileSAM 试图用 TinyViT 轻量级架构替换图像编码器的重量组件。

但是，它只通过与教师网络的图像嵌入的解耦知识蒸馏策略，使用紧凑的图像编码器网络进行训练。这种部分训练策略不可避免地导致性能衰减，而没有最终Mask预测的监督。FastSAM 将分割 anything 任务转移到具有一个 foreground 类别的实例分割任务，与 Yolov8 配合使用。为了实现可提示的分割，FastSAM 应用了一种后处理策略与实例分割网络相结合。然而，这种改写的框架无法在下游零样本任务上达到与 SAM  comparable 的性能。

为了进一步推动高效分割 anything 模型的极限，本文提出了一种完整的框架来获得 TinySAM，在降低计算成本的同时，尽可能地保持零样本分割能力。具体来说，作者提出了一种全阶段的知识蒸馏方法，以提高紧凑学生网络的能力。学生网络以端到端的方式，在教师网络不同阶段的监督下进行蒸馏。

此外，作者提出了一种在线硬提示采样策略，使蒸馏过程更加关注硬例子，从而提高最终性能。作者还将后训练量化适应到可提示的分割任务，并进一步降低计算成本。此外，作者发现，由于需要从网格提示点生成大量Mask，在图像上分割 everything 需要巨大的计算成本。

为此，作者提出了一种分层的分割 everything 策略，将 everything 推理加速 2 倍，几乎不降低性能。通过所提出的所有方法，作者的 TinySAM 导致了计算减少了数个数量级，并推动了高效分割 anything 任务的极限。例如，与原始 SAM 相比，TinySAM 可以实现 100 倍的加速。在各种零样本迁移任务上的广泛实验表明，作者的 TinySAM 与对应方法相比具有显著的优势。

2 Related Work

Segment Anything Model

最近提出的分割 anything 模型（SAM）在目标分割和下游视觉任务中证明了其通用性和灵活性。SAM 由三个子网络组成，即图像编码器、提示编码器和Mask解码器。图像编码器是一个基于重视觉 Transformer 的网络，它将输入图像提取成图像嵌入。提示编码器设计用于编码输入点、框、任意形状的Mask和自由形式文本，并使用位置信息。

几何提示和文本提示分别使用不同的网络进行处理。Mask解码器包含一个双向 Transformer ，将图像编码器和提示编码器的输出用于生成最终的Mask预测。与提出的 SA-1B 数据集相结合，该数据集包含 1.1 亿张高分辨率图像和超过 10 亿个高质量分割Mask，SAM 展示了针对任何类别和形状目标的令人印象深刻的高质量分割能力。

此外，SAM 在零样本下游视觉任务上展示了强大的泛化能力，包括边缘检测、目标 Proposal 、实例分割和文本到Mask预测。由于灵活的提示模式和高质量分割能力，SAM 被认为是视觉应用的基础模型。然而，SAM，尤其是图像编码器网络，包含大量参数，需要高计算能力进行部署。因此，在资源受限的边缘设备上应用 SAM 并不容易。SAM 的压缩和加速已成为一个重要的研究主题[49, 50]。

Knowledge Distillation

Hinton等人提出了一种知识蒸馏方法，通过教师网络的输出来监督轻量级学生网络的训练。自那时以来，知识蒸馏已成为在训练过程中改进紧凑网络性能的重要方法。知识蒸馏方法可以大致分为两类，即针对网络输出的蒸馏和针对中间特征的蒸馏。大多数关于知识蒸馏方法的研究都集中在图像分类任务上。后续的工作提出了用于高级计算机视觉任务（如目标检测和语义分割）的知识蒸馏方法。Zhang等人提出使用蒸馏方法来获得一个高效的分割 anything 模型（MobileSAM）。

然而，MobileSAM只使用图像和原始 SAM 的相应图像和图像嵌入来监督图像编码器网络。这种部分蒸馏策略可能导致轻量级学生网络从教师网络或Token数据中没有指导，无法获得Mask Level 信息，从而造成相当大的性能下降。

Quantization

模型量化也是常用的模型压缩方法之一，它将更高位宽的重量或激活值量化到较低位宽，以减少存储要求和计算复杂性，同时允许有限的准确性损失。模型量化方法可以分为两类，量化感知训练（QAT）和后训练量化（PTQ）。

QAT 方法需要一个带有标签的训练数据集和大量的训练成本，而 PTQ 方法只需要一个小型的无标签校准数据集，因此更加高效。许多先前的 PTQ 方法[30, 38]已经提出，以寻找卷积神经网络（CNN）的适当量化参数。

随着视觉 Transformer （ViT）在各种视觉任务上取得了显著的性能，最近的工作研究了如何将后训练量化应用于基于 ViT 的模型，并已使用 8 位量化配置实现了强大的性能。然而，对于提示性分割任务，尤其是分割 anything 模型，量化尚未得到充分探索。

3 Methodology

Overview of TinySAM

本文提出了一种框架，以获取高度高效的 SAM，如图 1 所示。

首先，在 3.2 节中，作者引入了一种专门为 SAM 设计的全阶段知识蒸馏。为了进一步激活蒸馏过程，作者使用了在线硬提示采样，以挖掘教师网络到学生网络之间的硬知识。其次，在 3.3 节中描述的后训练量化方法被适应到可提示分割任务，并应用于轻量级学生网络。第三，在 3.4 节中描述的分层 everything 推理模式被设计为分割 anything 任务，可以避免只有微小精度损失的巨额冗余计算，并加快推理时间，实现 的加速。

Full-Stage Knowledge Distillation

SAM 包括三个子网络，即图像编码器、提示编码器和Mask解码器。图像编码器网络基于视觉 Transformer，消耗大量的计算成本。受到 MobileSAM 的启发，作者使用轻量级 TinyViT 替换原始的沉重图像编码器网络。这种简单的替换存在相当大的性能衰减。因此，作者提出了一种全阶段知识蒸馏策略，在多个知识 Level 上指导轻量级图像编码器在学习过程中。

除了预测结果与真实标签之间的传统损失之外，作者在图 2 中引入了多个在不同阶段的蒸馏损失。

具体而言, 作者从教师网络的多个节点中选择来指导学生网络从多个知识 Level 的学习。首先, 作者将图像编码器的输出特征, 即图像嵌入, 作为蒸馏信息。图像嵌入集中了输入图像的信息, 这在预测过程中是基本知识。对于一个输入图像 , 图像嵌入的蒸馏损失函数可以表示为:

其中 和 分别表示学生网络和教师网络的图像编码器。由于图像 Level 的信息与 Mask 预测不直接相关, 对于这个分割任务, 与最终输出更接近的特征是必不可少的。自然地, 选择教师网络的最终输出作为蒸馏点。输出蒸馏损失 可以描述为:

其中 和 分别表示学生网络和教师网络的Mask解码器。 表示Mask解码器的 Query, 它是提示嵌入和输出Token的 ConCat 。由于 SAM 的结构相当复杂, 之前提到的两个蒸馏可能不一致且难以让轻量级学生网络学习。

作者进一步提出了一种将Mask解码器双向 Transformer 中的输出Token进行蒸馏的方法, 它从提示嵌入和图像嵌入中交互信息。它以更抽象的方式捕获目标Mask信息。相应的蒸馏损失 可以描述为:

其中 和 分别表示Mask解码器的双向 Transformer 模块, 表示损失函数。作者实际发现, 特征差值的数值可能使通常使用的 MSE 损失 ( 距离) 太小而无法良好地优化。因此,作者使用 距离函数代替。总的蒸馏损失函数 可以表示为:

其中 分别表示每个蒸馏损失的超参数。总的训练损失是 mask 预测 和 IoU 预测 的线性组合, 其中 是Focal Loss [29] 和 dice 损失 [37] 的组合, 是预测 IoUs 和计算 IoUs 之间的 L1 损失函数。

通常, 从带有标签的训练数据中随机采样可以用于生成驱动基于提示的Mask预测网络端到端训练的提示。为了进一步简化教师网络和轻量级学生网络之间的蒸馏的学习过程, 作者提出了一种硬提示采样策略, 使训练样本集中在预测的困难区域。以点提示为例, 最初在带有标签的 Mask区域 内部随机采样点 。这些初始点与输入图像一起输入到网络中, 得到预测的 Mask区域 。

然后, 作者从 和 的差集随机采样提示点, 并交互地进行这个过程。第 轮的采样点 来自 和 的差集, 即:

其中,

在应用于训练过程时, 第 次迭代从 0 到 9 随机采样, 这使得所采样提示的难度在一个受限的范围内。

Quantization

量化旨在将浮点张量 映射到 -位整数张量 并乘以缩放因子 。均匀对称量化可以如下表示:

对于矩阵乘法 , 它可以使用两个缩放因子 和 进行量化, 并将其表示为 。衡量 和 之间距离的度量对于优化 和 是至关重要的。遵循图像分类模型中量化方法的成功实践, 作者采用 Hessian 引导的度量方法作为距离来解决缩放因子, 这更符合任务损失。与分类任务不同, SAM 可提示分割任务的输出是包含细粒度Mask的分割预测。

因此, 作者使用Mask和 IoU 的 Kullback-Leib (KL) 散度作为任务损失, 并使用一些校准数据计算 Hessian 矩阵, 任务损失可以表示为:

在上述公式中, 和 分别表示浮点模型输出的预测结果, 和 表示经过量化后的输出。

在确定了距离度量之后, 作者可以将 和 作为交替迭代网格搜索问题来解决。通过校准数据, 作者可以得到 和 的最大值, 即 和 , 并使用两个参数 和 分别指定 和 的搜索范围, 即 和 . 这两个搜索范围分别线性划分为 个候选选项。 和 分别对几个迭代轮次进行交替优化。

在矩阵乘法之后经过softmax，输入在区间[0,1]的两端分布不均匀，而GELU 后的特征在正负范围内变化很大。这两种情况都与均匀量化的假设相去甚远，即神经网络中的激活遵循高斯分布。违反这一假设将导致高量化误差。因此，作者将特征分为两组，并使用两个缩放因子来减少量化误差。

Hierarchical Segmenting Everything

SAM 提出了一个自动Mask生成器, 以网格方式采样点来分割 everything。然而, 作者发现密集点网格会导致过细粒度的分割结果, 并占用大量的计算资源。一方面, 对于完整的目标, 过多的采样点可能会导致目标的某些不同部分被错误地分割为独立的Mask。另一方面, 由于图像编码器已经大幅缩减, everything 模式推理的时间成本主要在Mask解码器部分。

对于 SAM 自动Mask生成器的默认设置, 它将提示采样 个点, 这意味着Mask 解码器被推理 1024 次。在单个 V100 GPU 上, 图像编码器耗时 16ms, Mask解码器耗时 894 。

为了减少 everything 模式的时间成本, 作者提出了一个分层的Mask生成方法。如图 3 所示,作者的分层策略与原始策略进行了比较。与原始 everything 模式不同, 在第一步中, 作者只使用了每边 的点, 因此总点数为原始设置的 , 如图 3(c) 所示。然后, 作者使用这些提示来推理提示编码器和Mask解码器, 并得到结果。

然后作者笁选出置信度超过阈值 的Mask, 并将对应的位置Token为可能被视为最终预测的区域。对于这些区域, 由于它们被认为是置信度较高的实例的分割结果, 因此不需要再生点提示。然后作者以与原始设置相同但忽略上述区域的密度采样点。如图 3(d) 所示, 作者的策略忽略了第一头牛的草和身体上的大多数点。与此同时, 第二头牛和天空上的点被保留以进一步分割。

具体而言, 第二头牛在第一轮中被错误地分割为与第一头牛相同的目标。这种策略可以避免推理时间的成本和目标的过度细分。然后作者利用第二轮采样的点提示来获得Mask预测。最后, 将这两轮的结果合并并进行后处理以获得最终Mask。作者的方法忽略了超过 的点,从而带来了巨大的延迟减少。

4 Experiments

Implementation Details

作者使用 TinyViT-5M 作为轻量级学生图像编码器, 并采用 SAM-H 作为教师模型, 遵循先前的研究工作。作者只使用 SA-1B 数据集的 , 即 11000 张图像作为全阶段蒸馏的训练数据。作者采用 Adam 优化器, 并在 8 个 epoch 中训练学生网络。对于每个迭代, 作者根据在线硬提示采样策略 [3.2] 按 64 个提示采样。

为了加速蒸馏过程, 作者在训练时间内预先计算并存储了教师网络的图像嵌入。因此, 教师网络的图像编码器在训练时间内不需要重复计算。对于后训练量化, 作者量化所有卷积层、线性层、反卷积层和矩阵乘法层。在卷积层和反卷积层中, 作者使用通道缩放因子。而对于线性层和矩阵乘法层, 作者分别应用层缩放因子和头缩放因子。

对于迭代搜索, 作者设置 , rounds 。作者在 SA-1B 数据集上使用 8 张图像校准量化模型。作者在实例分割和点提示分割等下游任务上进行零样本评估。遵循 SAM 的建议, 作者采用多输出模式, 并认为最终Mask预测具有最高的 IoU 预测。

Zero-Shot Instance Segmentation

对于零样本实例分割任务，作者严格遵循 SAM 的实验设置，并使用 ViTDet-H 的目标检测结果作为实例分割的边界提示。作者在 COCO 数据集和 LVIS v1 数据集的基准上评估了零样本实例分割任务，并与不同的 SAM 变体以及先前的有效模型（如 FastSAM 和 MobileSAM）进行了比较。

如表 1 所示，与先前的方法相比，作者提出的 TinySAM 获得了优越的性能。具体而言，作者的 TinySAM 在 FLOPs 和实例分割精度方面超过了 FastSAM，即仅需 67% 的 FLOPs，即可实现约 4% AP 提升。

在相同的计算成本下，作者的 TinySAM 在 COCO 数据集上比 MobileSAM 实现了 0.9%+ 的 AP，而在 LVIS v1 数据集上实现了 1.6%+ 的 AP。作者的 W8A8 量化变体 TinySAM（Q-TinySAM）在不同方法之间的性能也具有竞争力。

具体而言，Q-TinySAM 在 COCO 数据集上比 MobileSAM [49] 实现了 0.3%+ 的 AP，仅需 26% 的 FLOPs。作者在 COCO 验证集上的部分视觉结果如图 5 所示。与 SAM 变体的其他有效变体相比，作者提出的 TinySAM 捕获了更清晰的边界。

Zero-shot Points Valid Mask Evaluation

在这个部分, 作者评估了作者的 TinySAM 在使用多个点作为提示分割物体的性能。作者使用与先前工作相同的点选择指标, 该指标计算假阳性和假阴性Mask的距离变换, 然后以最大值采样点。作者计算每个数据集的 mIoU 来评估不同模型的性能。

作者选择在SAM中使用的 23 个数据集的子集进行高效的评估, 这些数据集包括 BBBC038v1, DOORS, TimberSeg 和 LVIS。为了进行公平的比较, 作者遵循 SAM 论文中 Table 7 的设置来选择图像和Mask, 并使用相应分裂中的前 个Mask进行评估。

在这个部分, 作者评估了作者的 TinySAM 在使用多个点作为提示分割物体的性能。作者使用与先前工作相同的点选择指标, 该指标计算假阳性和假阴性Mask的距离变换, 然后以最大值采样点。作者计算每个数据集的 mIoU 来评估不同模型的性能。

作者选择在SAM中使用的 23 个数据集的子集进行高效的评估, 这些数据集包括 BBBC038v1, DOORS, TimberSeg 和 LVIS。为了进行公平的比较, 作者遵循 SAM 论文中 Table 7 的设置来选择图像和Mask, 并使用相应分裂中的前 个Mask进行评估。

Everything Mode Acceleration

作者在 COCO 验证集上评估作者提出的分层 everything 推理策略。所有延迟基准都在单个 NVIDIA V100 GPU 上进行。作者从 val2017 中采样 100 张图像，并在这些样本上进行 everything 模式推理。everything 模式中的阈值都保持与默认值相同。结果如表 2 所示。作者对使用不同策略评估的相同模型应用相同的阈值和稳定性分数，以进行公平比较。

但是, 这些模型之间的结果可能不同。作者的分层策略在与原始 点网格策略相当的情况下, 实现了可比较的结果, 同时推理时间成本降低了约 。

图 6 显示了点网格策略与所提出分层策略之间的视觉比较。作者可以发现，所提出的分层策略在某些情况下避免了目标的过度细分，从而弥补了 IoU 的一些损失。更高效的 everything 模式推理的更多视觉结果如图 8 所示。

Ablation Studies

在本节中，作者在 COCO 验证数据集上进行零样本实例分割任务的 ablation 研究。实验设置与第 4.2 节相同。

不同模块的影响。作者首先分别评估了不同模块的影响，即全阶段知识蒸馏损失、在线硬提示采样和后量化。

如表 3 所示, 采用作者提出的全阶段知识蒸馏策略, 性能从 提高到 。进一步与在线硬提示采样策略相结合, 作者的方法可以获得另一个 增益。使用后训练量化导致 降级, 但同时也大大降低了计算成本, 这在主要的实验中也有所讨论。

如图 7 所示，作者展示了为全阶段知识蒸馏设计的硬提示采样的采样过程。星号表示具有不同迭代次数的采样点。经过一些迭代，采样区域更接近目标Mask的边缘，这使得提示相对更难让学生网络学习。

不同蒸馏损失的影响。对于详细的全阶段知识蒸馏过程，作者研究了从教师网络进行三层蒸馏的必要性。表 4 显示了不同蒸馏损失组合的 ablation 结果。输出蒸馏损失起着重要作用，因为它接近监督信息，并且与教师网络的相似性直接反映在评估指标中。Token损失和嵌入损失都被证明是有益的，因为它们与教师网络的关键节点相关，分别反映了图像 Level 的信息和提示与图像之间的交互。

不同的量化位数。作者在这里探索不同的量化位数的影响。表 5 报告了在 COCO 数据集上的平均精确度。从结果中，作者可以得出结论，将量化位数提高到 8 位只会导致轻微的性能下降。作者还通过进一步减少量化位宽到 6 位来演示性能。

5 Conclusion

在本文中，作者提出了一种框架，以推动分割 anything 任务的极限并获取一个高度高效的模型，命名为 TinySAM。作者首先提出了一种全阶段知识蒸馏方法，采用在线硬提示采样策略来蒸馏一个轻量级的学生模型。作者还适应了后训练量化到可提示的分割任务，并进一步降低了计算成本。

此外，作者提出了一种分层的分割 everything 策略，可以将 everything 推理加速 2 倍，几乎不降低性能。通过所提出的所有方法，作者的 TinySAM 导致了计算减少了数个数量级，并推动了高效分割 anything 任务的极限。在各种零样本迁移任务上的广泛实验表明，作者的 TinySAM 与对应方法相比具有显著的优势。作者希望能够提出的 TinySAM 为设计一个高度高效的分割 anything 模型带来有益的视角。

Appendix

作者在附录中提供了更多的可视化结果。

图 9 展示了在 LVIS v1 数据集上的零样本实例分割。为了清晰地呈现，只显示 VitDet-H 分数高于 0.8 的检测框。LVIS 数据集的标签比 COCO 数据集更细致，在考虑准确性和效率方面，所提出的 TinySAM 显示了更大的优势。

图 10 显示了通过提出的 TinySAM 模型进行分层 everything 推理和其对应算法的结果。TinySAM 捕获了清晰的边界并产生了更细粒度的Mask，而 MobileSAM 和 FastSAM 有时会产生人造的边界和Mask。TinySAM 显示了更接近原始 SAM 的性能，同时消耗了显著更少的计算成本。

参考

[1].  TinySAM: Pushing the Envelope for Efficient Segment Anything Model

公众号后台回复“数据集”获取100+深度学习各方向资源整理