[作者解读] SAM2-UNet: Segment Anything 2 Makes Strong Encoder for Natural and Medical Image Segmentation

[论文地址] [代码] [arXiv 24.08]

Abstract

图像分割在视觉理解中发挥着重要作用。最近，新兴的视觉基础模型在各种任务中不断取得优异的性能。继这些成功之后，我们在本文中证明了 Segment Anything Model 2 (SAM2) 可以成为 U 型分割模型的强大编码器。我们为多功能图像分割提出了一个简单而有效的框架，称为 SAM2-UNet。具体来说，SAM2-UNet 采用 SAM2 的 Hiera 骨干作为编码器，而解码器则采用经典的 U 型设计。此外，还在编码器中插入了Adapter，以便进行参数高效微调。在各种下游任务(如伪装目标检测、显著目标检测、海洋动物分割、镜面检测和息肉分割)上进行的初步实验表明，我们的 SAM2-UNet 完全可以击败现有的最先进专业方法，而无需任何花哨设计。

---

TL;DR

SAM2的Hiera Backbone可以直接单独拿出来给其他任务用，效果有明显提升。(即标题中的"Makes Strong Encoder")

---

Framework

SAM2-UNet的网络结构如下所示：

整体由四部分构成：

• Encoder：采用的是SAM2中的Hiera Backbone。由于Hiera同样具有层次化结构，因此其完全可以用于替换一些现有的经典Encoder，例如ResNet，Swin，PVT等等。
• RFBs：用于降采样兼提取多尺度特征。该部分理论上可以使用任意其他Attention模块替换。更进一步的，也可以直接放弃降维直接将原始Encoder特征跨层连接至Decoder。
• Adapter：用于实现参数高效微调，以降低训练的显存占用。参数高效微调的原理可以参见此处。
• Decoder：采用了标准的U-Net Decoder设计，每个Decoder Block包括两个Conv-BN-ReLU块。

---

Result

由于SAM2本身是在分割任务上预训练，因此finetune至下游分割任务时效果明显好于基于ImageNet预训练的方法。以伪装目标检测的性能为例，对比了SINet (CVPR 20), PFNet (CVPR 21)，ZoomNet (CVPR 22)，FEDER (CVPR 23)：

该结果甚至能打一些基于SAM1的方法，例如DSAM (ACM MM 24)。

---

Ablation

主要讨论了不同尺寸的Hiera。哪怕对于Hiera-Tiny，其结果依然能打过一些较老的方法如SINet和PFNet。而Hiera-Tiny的参数量为28M，对比现有的一些backbone，如Res2Net-50的25.7M，消耗增加并不算多，因此完全是有可能在后续工作中进行替换的。

---

Q&A

Q1： 训练的显存占用情况？
A1： 在4090机器上，batch size 12占用显存约为16G，如果调小batch size完全能在2080ti等小显存机器上运行。

Q2： 输入分辨率为什么是352，而非SAM2默认的1024？
A2： 尝试过调整为1024，有的数据集性能会略微上升，而有的变化不大。由于网络本身的decoder设计非常简单，并没有做大感受野设计，因此直接提升分辨率容易导致全局特征出现问题。此外，使用1024分辨率也会明显地降低训练速度。

Q3： SAM1的Vanilla ViT是否也能用类似的思路？
A3： 之前做过一个类似的工作Mammo-SAM，将SAM1的decoder去掉换成U-Net风格的级联解码器，但效果有限，而且引入U-Net中经典的Skip Connection后更难训练，不像SAM2这样直接方便。不过这几天也有同期工作讨论了对应的做法，例如SAM-UNet。

Q4： 与其他SAM2微调方法的对比？
A4： 另一个经典工作为SAM2-Adapter。从分割性能来讲，两者互有胜负。从结构来讲，SAM2-UNet删去了SAM2中大部分的模块，上手难度会明显更低，后续可扩展上限更高；而SAM2-Adapter对SAM2的结构基本进行了保留，更适合需要使用SAM特性的场景(如Prompting)。

Q5: 后续可以改进的方向？
A5: 例如可以把decoder做的更复杂，改善Adapter的结构和插入位置等等。

---

Citation

如果SAM2-UNet对你有所启发或帮助，欢迎来我们的GitHub点一个Star以及引用本篇论文。非常感谢！

@article{sam2unet,
  title={SAM2-UNet: Segment Anything 2 Makes Strong Encoder for Natural and Medical Image Segmentation},
  author={Xiong, Xinyu and Wu, Zihuang and Tan, Shuangyi and Li, Wenxue and Tang, Feilong and Chen, Ying and Li, Siying and Ma, Jie and Li, Guanbin},
  journal={arXiv preprint arXiv:2408.08870},
  year={2024}
}