【读点论文】Group-Mix SAM: Lightweight Solution for Industrial Assembly Line Applications，MobileSAM做教师模型_group-mix sam(group-mix sam: lightweight solution -CSDN博客

本文链接：https://blog.csdn.net/weixin_43424450/article/details/140401664

Group-Mix SAM: Lightweight Solution for Industrial Assembly Line Applications

Abstract

自从大约一年前 Segment Anything 模型 (SAM) 问世以来，它引起了学术界的极大兴趣，并从各个角度催生了大量的研究和出版物。然而，由于 SAM 的图像编码器非常大，高达 632M，因此它尚未在实际的装配线场景中部署。在本研究中，我们用轻量级的图像编码器替换了重量级的图像编码器，从而使 SAM 能够在实际的装配线场景中部署。具体来说，我们采用解耦蒸馏在资源有限的环境中训练 MobileSAM 的编码器。整个知识蒸馏实验可以在单个 RTX 4090 上在一天内完成。与 MobileSAM 相比，由此产生的轻量级 SAM（称为 Group-Mix SAM）的参数减少了 37.63%（2.16M），浮点运算减少了 42.5%（15614.7M）。然而，在我们构建的工业数据集 MALSD 上，其 mIoU 仅略低于 MobileSAM，为 0.615。最后，我们进行了全面的对比实验，以证明 Group-Mix SAM 在工业领域的优越性。凭借其出色的性能，我们的 Group-Mix SAM 更适合实际的流水线应用。
liliu-avril/Awesome-Segment-Anything：此存储库用于对 Meta AI 的 Segment Anything 模型（SAM）进行首次全面调查。 (github.com)
论文地址：[2403.10053] Group-Mix SAM: Lightweight Solution for Industrial Assembly Line Applications (arxiv.org)

Introduction

ChatGPT彻底改变了NLP领域，在大规模模型上取得了突破性进展。在计算机视觉领域，SAM的作用无疑是突出的。SAM在零样本实例分割任务中取得了最先进的性能，取得了巨大的成功。SAM与其他模型兼容，可以实现文本引导分割、细粒度图像编辑等高级视觉应用。尽管有这些优势，但SAM模型很难部署在实际的流水线上。原因是在实际部署中使用SAM执行分割任务会导致难以承受的计算和内存成本。SAM可以看作两个部分：基于ViT的图像编码器和提示引导的掩码解码器。基于ViT的图像编码器的参数大小比提示引导的掩码解码器大两个数量级，大小分别为632M和3.87M。掩码解码器是轻量级的，所以SAM无法部署在实际流水线上的根本原因是图像编码器的参数大小。
在实际的流水线中，负责运行算法的是边缘计算机，这些边缘计算机的购买通常受限于价格，存在内存不足、计算能力弱等问题。因此，为了在实际的流水线场景中部署使用，我们需要将重量级的图像编码器替换为轻量级的图像编码器，以减小 SAM 的大小。
本文针对MobileSAM因边缘计算机计算能力不足、内存占用过多而无法部署在流水线上的问题，提出了用更小的图像编码器结构Groupmixformer取代MobileSAM编码器中原有的ViT-T结构，取得了很好的效果，将得到的模型命名为Group-Mix SAM。如[Segment anything]所述，在256个A100 GPU上用ViT-H图像编码器训练SAM需要68小时。直接用小型图像编码器训练费时费力，效果也不一定好。因此，我们转向知识蒸馏技术。解耦蒸馏不依赖组合解码器，而是直接从原 SAM 的 ViT-H 中提取出小图像编码器，与半耦合（冻结掩码解码器，从掩码层优化图像编码器）和耦合蒸馏（直接从掩码层优化图像编码器）[Mobilesam] 相比，在时间和效率上均具有优势。因此，我们选择解耦蒸馏作为知识蒸馏的方法。具体来说，它将知识从基于 ViT-T 的结构转移到具有较小图像编码器的 Groupmixformer 中。

Related works

Vision Transformers

Vision Transformers (ViTs) 通过使用多头自注意力 (MHSA) 机制，在视觉应用中表现出惊人的性能，有效地捕捉了全局依赖关系。与 CNN 相比，ViT 模型表现出更高的上限和泛化能力，但由于参数数量较多，也存在实际部署问题。因此，许多关于 ViT 的论文和项目都专注于高效部署，例如 LeViT、MobileViT、Tiny-ViT、EfficientViT、GroupMixFormer 等。设计高效的 ViT 结构可以构建高效的 SAM。

Segment Anything Model

SAM 广泛应用于视觉任务，因为它可以根据点或框提示分割图像中的任何对象。它的零样本泛化是基于 SA-1B 数据集学习的，并且在各种视觉任务中表现出色。此外，SAM 还应用于医学图像分割和透明物体检测等领域。由于其广泛的实际应用，SAM 的实际生产线部署受到越来越多的关注。最近的研究提出了一些降低 SAM 计算成本的策略，例如 EfficientSAM，它引入了一个名为 SAMI 的预训练框架，该框架利用了 SAM 的掩模图像；EdgeSAM，它将原始基于 ViT 的 SAM 图像编码器提取为纯基于 CNN 的架构；MobileSAM，它通过解耦蒸馏获得轻量级图像编码器来取代 SAM 原始笨重的编码器。我们的重点是实现 SAM 的实际生产线部署。

Knowledge Distillation

知识蒸馏（KD）是一种用于迁移模型能力的方法。根据迁移方法的不同，可以分为两个主要方向：目标蒸馏（logits方法蒸馏）和特征蒸馏。目标蒸馏中最经典的论文是[Distilling the knowledge in a neural network]，其中提出了软目标和温度等概念。而特征蒸馏则是让学生学习老师网络结构的中间层特征。最早提出该方法的论文是FitNets，它强制学生的中间层近似老师相应的中间层特征，从而达到特征蒸馏的目的。知识蒸馏的大部分研究集中在分类、语义分割和物体检测等任务上。

Methodology

本节将介绍论文中使用的几种关键方法，首先介绍我们最终选定的图像编码器：Groupmixformer及其原理和优势，然后介绍我们采用的蒸馏方法：解耦蒸馏及其实现原理。

Groupmixformer

Groupmixformer 提出了 Group-Mix Attention (GMA) 作为传统自注意力机制的高级替代方案。与流行的多头自注意力机制（仅对每个 token 之间的相关性进行建模）相比，GMA 可以利用组聚合器同时捕获 token-to-token、token-to-group 和 group-to-group 之间的相关性。Groupmixformer 在图像分类、对象检测和语义分割等下游任务中取得了优异的表现。在4.1节中，我们对适用于工业生产线场景的Encoder进行了充分的对比实验，其中Groupmixformer表现最优。因此，我们选取其作为知识提炼的学生模型。

Distillation

在知识蒸馏中，教师和学生架构是知识传递的载体，也就是说，教师向学生传递知识的质量取决于教师和学生的网络结构如何设计。根据之前的经验，我们选择了解耦的蒸馏方法，使用 MobileSAM 中的编码器 vit-t 作为教师模型，Groupmixformer 作为学生模型。蒸馏框架我们采用了 nanosam [https://github.com/NVIDIA-AI-IOT/nanosam]，它使用 NVIDIA TensorRT 大大加速了蒸馏过程中教师模型的推理速度。此外，它使用教师和学生通过前向推理得到的特征图来计算最终的损失，这种简单的机制也加速了蒸馏过程。因此，我们完整的蒸馏过程可以低成本地在不到一天的时间内复现到单台 RTX 4090 上

Experiments and Analysis

本节将介绍本文的实验过程、结果和分析。首先介绍选择适合实际工业流水线场景的编码器的实验过程和结果。然后介绍知识蒸馏所使用的数据集和实验设置，并展示蒸馏实验的结果。最后比较 Group-Mix SAM 和 MobileSAM 的性能。

Experiments on selecting image encoders

我们需要选择适合实际工业流水线场景的编码器。为此，我们通过收集美的实际流水线视频数据创建了美的流水线场景数据集 (MALSD)。如下图所示，MALSD 只有两个类别：Act0 和 NG。Act0 指的是工人手持产品说明书的图像，而 ng 指的是其他杂项图像。训练集分别有 533 张和 668 张 Act0 和 NG 图像，测试集分别有 144 张和 238 张 Act0 和 NG 图像。
- 美的装配线场景数据集（MALSD）示例
如下表所示，我们在 MALSD 上训练和测试了现有的最先进图像编码器。我们观察到 groupmixformer 在 MALSD 上的表现最佳。它在验证集上的表现比第二好的模型 LeViT 高出 3.28%。因此，我们得出结论，Groupmixformer 更适合我们的工业场景。
- 在美的装配线数据集上训练和测试现有的图像编码器。表中的粗体表示最佳，下划线表示第二佳。

Experimental setup for knowledge distillation

我们选择了 Nanosam 作为蒸馏框架。我们使用 COCO 2017 数据集，其中未标记图像用于训练，标记图像用于测试。蒸馏期间输入给学生的图像大小为 1024，批处理大小为 8，学习率为 3e-4，所有实验统一设置 13 个 epoch。使用的损失函数是 Huber。为了减少蒸馏所需的时间，我们使用 TensorRT 加速了教师模型。因此，整个蒸馏过程在不到一天的时间内就在单个 RTX 4090 上完成。

Experimental results of knowledge distillation

结果如下表所示。我们选择了MobileSAM的编码器Vit-T作为教师模型，其参数量为5.74M，浮点运算量为36742.79M，在coco数据集上准确率为72.8%，在小模型中也表现出色。我们采用了三种不同的结构，Resnet18、与教师模型结构相同的VitT、Groupmixfromer（GMF）作为学生模型进行知识蒸馏。随后，我们测试了Groupmixfromer的内部结构，例如将stage数从4个减少到3个，降低层的串行深度，调整学习率等。最终，我们确定了末行结构为 [2,2,8,2] 的Groupmixfromer最适合我们的工业场景。虽然准确率比Vit-T低0.3%，但参数减少了37.63%（2.16M），浮点运算减少了42.5%（15614.7M），非常适合工业流水线场景。训练过程中的特征对比图如下图所示。
- 知识蒸馏过程中老师和学生的特征对比。图中数字代表epoch，特征大小为(1,256,64,64)，这里可视化第一维(64*64)。

Performance of Group-Mix SAM

在图像分割中使用经过提炼、适合工业流水线场景的图像编码器 GMF，得到了 Group-Mix SAM。我们在 COCO 公开数据集上定性比较了经过提炼的 Teacher 模型、GMF、Vit-T 和 Resnet18 主干模型（它们取代了 Mobile SAM 中的编码器）。我们在实验中使用边界框作为提示。如下图所示，Group-Mix SAM 似乎在某些细节上表现更好，例如第一排和第三排的马尾和背部，而 Teacher 模型、Resnet18 和 Vit-T 都存在一些缺陷。
- 教师模型和学生模型（即 GMF、Resnet18 和 Vit-T）在 coco 数据集上的性能比较。建议以彩色和放大的方式查看比较结果，以获得更好的可视化效果。
我们重复上述实验，并在我们自己构建的数据集 MALSD 上进行了定性比较，这次以工人手部中心点为提示。如下图所示，尽管参数和浮点运算次数最少， GMF 仍然表现良好。例如，在第一行和第四行工人手持说明书，以及第二行工人手持条码扫描器时，GMF 与其他三种方法相比能够准确、清晰地分割出手部。
- 教师模型和学生模型（即 GMF、Resnet18 和 Vit-T）在 MALSD 数据集上的性能比较。建议以彩色和放大的方式查看比较结果，以获得更好的可视化效果。
最后，我们利用上述四种方法在MALSD上对Group-Mix SAM进行了详细测试，结果如下表所示。可以看出，在MALSD中，mIoU比常规数据集测量的要低（MobileSAM在常规数据集上可以达到70%以上），这是因为工业数据集中干扰和噪声较多。虽然Group-Mix SAM使用的编码器是从MobileSAM中提炼出来的（参数减少了37.63%（2.16M），浮点运算减少了42.5%（15614.7M），但在工业数据集上差别不大。
- mIoU 比较。假设来自原始 SAM 的预测掩码是基本事实，则较高的 mIoU 表示更好的性能。

Conclusion

我们从自建工业数据集MALSD中挑选出适合实际工业流水线场景的编码器，然后通过解耦蒸馏，将MobileSAM的编码器蒸馏到Gounpmixformer架构中，通过用轻量级图像编码器替换重量级图像编码器，使SAM能够部署到实际的流水线场景中，凭借其优异的性能，我们的Group-Mix SAM更适合实际的流水线应用。