使用 AMD GPU 实现 Segment Anything

Segment Anything with AMD GPUs — ROCm Blogs

作者： Sean Song

发布日期：2024年6月4日

介绍

分割任务——识别图像中哪些像素属于某对象——是计算机视觉中的一个基础任务，应用广泛，从科学图像分析到照片编辑。Segment Anything 模型(SAM)是一个先进的图像分割模型，它通过提示分割（promptable segmentation）实现了前所未有的多功能性，使图像分析任务变得更加简单。SAM 可以用于帮助在需要查找和分割图像中任何对象的领域内应用。对于AI研究社区和其他相关领域，SAM 很有可能成为大型AI系统中的关键组件，在多模态环境中实现对世界的全面理解。

在这篇博客中，我们将演示如何在使用 ROCm 的 AMD GPU 上运行 Segment Anything 模型。

SAM

SAM经过训练可以根据各种提示返回有效的**分割掩码**，这些提示包括前景和背景**点**、大致的**框**或**掩码**、非结构化的**文本**，或任何其他指示图像中要分割内容的标识。对有效掩码的要求仅仅意味着即使提示是模糊的，可能指向多个物体（例如，衣服上的一个点可能表示衣服或穿衣服的人），输出结果也应是这些物体中的一个合理掩码。

SAM的先进设计使其能够在没有事先知识的情况下适应新的图像分布和任务——这种特性被称为零样本迁移。SAM是在庞大的SA-1B 数据集上进行训练的，该数据集包含超过10亿个掩码，分布在1100多万张精心挑选的图像上（见下方来自SAM论文的示例），已经展示了令人印象深刻的零样本性能，在许多情况下超越了之前的全监督结果。

图片来源：Segment Anything（Kirillov et al.）。

图片来源：Introducing Segment Anything (Meta Research)。

SAM 模型由三个关键模块组成:

• 图像编码器: 为了实现可扩展性和强大的预训练方法，SAM 采用经过 Masked AutoEncoder (MAE) 预训练的 Vision Transformer (ViT)，经过最小化调整后以处理输入图像。图像编码器每个图像只运行一次，可以在提示模型之前应用。

• 提示编码器: SAM 考虑了两组提示：稀疏提示（点、框、文本）和密集提示（掩码）。它通过位置编码加上每种提示类型的学习嵌入来表示点和框，并使用来自 CLIP 的开箱即用文本编码器对自由文本进行编码。密集提示（即掩码）使用卷积进行嵌入，并与图像嵌入按元素相加。

• 掩码解码器: 掩码解码器高效地将图像嵌入、提示嵌入和输出令牌映射到掩码。这一设计采用修改后的 transformer 解码器块，随后是动态掩码预测头。修改后的解码器块使用提示自注意力和交叉注意力（从提示到图像嵌入和反之亦然）来更新所有嵌入。之后，模型对图像嵌入进行上采样，并通过多层感知机（MLP）将输出令牌映射到动态线性分类器，然后计算每个图像位置的掩码前景概率。

接下来，我们将在三个部分中展示 Segment Anything 模型在具有 ROCm 的 AMD GPU 上的流畅执行：

• 自动生成所有对象的掩码

• 使用点作为提示生成掩码

• 使用框作为提示生成掩码

注意：在撰写此博客时，尽管 SAM 论文中探讨了文本提示的功能，但此功能尚未完全发布。本文不涉及文本提示。

我们在本博客文章中引用了 SAM 的 GitHub 代码仓库。可以在 facebookresearch/segment-anything 找到模型的源代码。

设置

本演示使用以下硬件和软件环境。有关全面的支持详细信息，请参阅 ROCm 文档。