SAM 模型架构

相关源文件

• segment_anything/modeling/image_encoder.py
• segment_anything/modeling/sam.py
• segment_anything/modeling/transformer.py
• segment_anything/utils/onnx.py

本文档提供了关于 Segment Anything Model (SAM) 架构的详细技术概述。重点关注模型本身的设计和内部组件。有关通过接口使用模型的信息，请参阅 SamPredictor 和 Automatic Mask Generator。

1. 概述

SAM 由三个主要组件组成，它们协同工作以执行分割。

来源：segment_anything/modeling/sam.py18-48

2. 模型组件

2.1 图像编码器

图像编码器是一个经过修改的 Vision Transformer (ViT)，它将输入图像处理成特征表示。

ImageEncoderViT 具有以下关键特征：

• Patch 嵌入：将图像分割成固定大小的 patches，并将它们投影到嵌入维度。
• 位置嵌入：可以使用绝对位置嵌入和相对位置嵌入。
• Transformer 块：一系列自注意力块，其中一些块具有窗口注意力以提高效率。
• 卷积颈：一个两层卷积网络，用于调整输出特征维度。

配置参数

参数	描述	默认
img_size	输入图像尺寸	1024
patch_size	图像 patches 的大小	16
embed_dim	嵌入维度	768
depth	Transformer 块的数量	12
num_heads	注意力头的数量	12
out_chans	输出特征通道数	256
window_size	局部注意力的窗口大小	0 (全局)

来源：segment_anything/modeling/image_encoder.py16-116

2.2 提示编码器

提示编码器将用户输入的各种提示（点、框或掩码）转换为可由掩码解码器使用的嵌入。

提示编码器处理三种类型的提示：

1. 点：每个点都与其位置和标签（前景/背景）一起被编码。
2. 框：框的角被编码为特殊点嵌入。
3. 掩码：先前的掩码通过卷积网络被编码为密集嵌入。

来源：segment_anything/modeling/sam.py106-110

2.3 掩码解码器

掩码解码器结合图像嵌入和提示嵌入来预测分割掩码。

掩码解码器的特点：

• TwoWayTransformer：处理稀疏和密集 token，具有交错的自注意力和交叉注意力。
• 掩码注意力：将转换后的特征应用于图像嵌入。
• 多输出：同时产生掩码 logits 和 IoU（质量）预测。

双向 Transformer 使得信息能够在稀疏提示 token 和密集图像 token 之间双向流动，从而有效地集成这些不同的模态。

来源：segment_anything/modeling/transformer.py16-240

3. 端到端处理流程

下图显示了 SAM 模型从输入到输出的完整数据流。

模型前向传播流程

1. 输入图像经过归一化、填充成方形，然后由 ImageEncoderViT 处理以生成嵌入。
2. 用户提示（点、框、掩码）由 PromptEncoder 编码成稀疏和密集嵌入。
3. MaskDecoder 使用 TwoWayTransformer 处理这些嵌入并预测掩码。
4. 低分辨率掩码 logits 被放大并进行阈值处理，以生成最终的二值掩码。

来源：segment_anything/modeling/sam.py54-174 segment_anything/utils/onnx.py107-144

4. 模型变体

SAM 有三种尺寸，取决于用于图像编码器的 ViT 主干。

变体	图像编码器	参数	特性
SAM-B	ViT-B	约 9100 万	最小的模型，推理速度最快
SAM-L	ViT-L	约 3.08 亿	中等尺寸，平衡性好
SAM-H	ViT-H	约 6.36 亿	最大的模型，最高质量的输出

每个变体使用相同的架构，但 ViT 主干的大小不同。有关模型尺寸考虑因素的更多信息，请参阅 使用不同模型尺寸。

来源：segment_anything/modeling/sam.py18-48

5. 技术实现细节

图像编码器细节

图像编码器使用 Vision Transformer 架构，其图像处理流程如下：

1. 图像被分割成不重叠的 patches（默认为 16x16 像素）。
2. 每个 patch 被线性投影到嵌入维度。
3. 添加位置嵌入（绝对或相对）。
4. 嵌入通过一系列带有自注意力的 Transformer 块。
5. 卷积“颈”将输出转换为所需的特征维度。

关键实现说明：

• 窗口注意力可选择性地用于降低大型图像的计算复杂度。
• 最终输出的形状为 (B, out_chans, H/patch_size, W/patch_size)。

来源：segment_anything/modeling/image_encoder.py119-182 segment_anything/modeling/image_encoder.py185-240

双向 Transformer 细节

双向 Transformer 是 SAM 架构的关键创新。

1. 每个 Transformer 块有四个组件：

   • 稀疏 token（提示）上的自注意力。
   • 从稀疏 token 到密集 token 的交叉注意力（提示 → 图像）。
   • 稀疏 token 的 MLP。
   • 从密集 token 到稀疏 token 的交叉注意力（图像 → 提示）。

2. 这种双向流动允许信息在提示和图像表示之间传递。

3. 最后的注意力层对 token 到图像的嵌入执行最终交互。

来源：segment_anything/modeling/transformer.py109-182