HRSAM：高分辨率图像上的高效交互式分割
计算机视觉与模式识别
2024-11-26 v2 人工智能
摘要
Segment Anything Model (SAM) 已推进交互式分割，但因高分辨率图像上的高计算成本而受到限制。这需要对图像进行下采样以满足 GPU 约束，牺牲了对高精度交互式分割所需的细粒度细节。为解决 SAM 的局限性，我们关注视觉长度外推问题，提出一种轻量级模型 HRSAM。外推使 HRSAM 在低分辨率上训练后能够泛化到高分辨率。我们首先发现外推与注意力得分之间的关联，从而将 HRSAM 基于 Swin 注意力机制。随后，我们引入 Flexible Local Attention (FLA) 框架，使用经过 CUDA 优化的高效记忆注意力加速 HRSAM。在 FLA 中，我们实现了 Flash Swin 注意力，相较于传统 Swin 注意力实现了 35% 以上的加速，并提出一种仅使用键值 (KV) 的填充机制来增强外推。我们还开发了 Cycle-scan 模块，该模块使用状态空间模型高效扩大 HRSAM 的感受野。进一步在 FLA 中开发 HRSAM++，通过引入锚点图提供外推的多尺度数据增强，并在稍微增加计算成本的情况下获得更大的感受野。实验表明，在标准训练下，HRSAM 的性能在延迟仅为前 SOTA 的 38% 时便超越了前 SOTA。通过 SAM 蒸馏，外推使 HRSAM 在更低的延迟下超过教师模型。进一步微调后，HRSAM 的性能显著超越了前 SOTA。
引用
@article{arxiv.2407.02109,
title = {HRSAM: Efficient Interactive Segmentation in High-Resolution Images},
author = {You Huang and Wenbin Lai and Jiayi Ji and Liujuan Cao and Shengchuan Zhang and Rongrong Ji},
journal= {arXiv preprint arXiv:2407.02109},
year = {2024}
}