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}
}