Ultralytics SAM 系列模型网络模块深度解析：blocks.py 中的注意力与特征融合构件

Ultralytics 仓库中的 ultralytics/models/sam/modules/blocks.py 是整个 SAM（Segment Anything Model）、SAM 2 及 SAM 3 系列架构共享的"模块库"，集中实现了视觉 Transformer 中最重要的可复用构件：随机深度正则化、mask 降采样、ConvNeXt 特征融合、双路 Transformer 解码器、RoPE 旋转位置编码、多尺度注意力、正弦/随机位置编码以及 patch 嵌入等。本文将结合仓库源码逐类拆解这些模块的职责、构造参数与数据流，帮助你建立从"网络层"到"完整分割模型"的清晰认知。

模块定位：blocks.py 在 SAM 架构中的角色

blocks.py 被多个上层模块直接引用，是观察 SAM 系列模型架构演进的绝佳窗口。从仓库源码可见其调用关系：

• encoders.py 中的 Hiera 图像编码器、Mask 编码器引用了 Block、CXBlock、Fuser、MaskDownSampler、MultiScaleBlock、PatchEmbed、PositionEmbeddingSine、PositionEmbeddingRandom；
• sam.py 中的 SAM 2 解码器使用 SAM2TwoWayTransformer 完成 image-query 双向往返注意力；
• memory_attention.py 中的视频记忆注意力使用 RoPEAttention 注入旋转位置编码；
• sam3/vitdet.py 复用 PatchEmbed 作为 ViT 检测骨干的 patch 化入口；
• build.py 则通过 TinyViT（见 tiny_encoder.py）构建轻量级图像编码器。

换言之，本文件聚合了"卷积类前处理构件 + Transformer 核心注意力构件 + 位置编码构件"三类基础单元，构成 SAM 系列图像/视频分割模型网络栈的基石。下文按功能分组逐类剖析。

一、训练正则化：DropPath 随机深度

DropPath 实现了训练阶段常用的 stochastic depth（随机深度）正则化，用于深层网络中随机丢弃部分残差分支，从而提升泛化能力。

class DropPath(nn.Module):
    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

    def forward(self, x):
        if self.drop_prob == 0.0 or not self.training:
            return x
        keep_prob = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
        if keep_prob > 0.0 and self.scale_by_keep:
            random_tensor.div_(keep_prob)
        return x * random_tensor

关键行为（对应 blocks.py 第 20-51 行）：

• 训练/推理解耦：drop_prob == 0 或处于 eval 模式时直接恒等返回，保证推理结果确定性；
• 批级掩码：掩码张量形状为 (batch, 1, 1, ...)，即同一 batch 内整条样本路径被整体保留或丢弃；
• keep-prob 缩放：scale_by_keep=True 时将保留值除以 keep_prob，等价于 inverted dropout，使训练期期望输出与推理期一致。

在文件内部，DropPath 被 CXBlock、MultiScaleBlock、Block 三个模块统一用作残差分支的门控（drop_path > 0.0 时启用，否则退化为 nn.Identity()）。

二、Mask 降采样与特征融合：MaskDownSampler、CXBlock、Fuser

1. MaskDownSampler：mask 的渐进式卷积编码

MaskDownSampler 以逐步下采样 + 通道扩张的方式，把单通道 mask 编码为 embed_dim 维特征（blocks.py 第 54-125 行）。

构造参数语义：

参数	默认值	说明
embed_dim	256	输出特征通道数
kernel_size	4	每层卷积核尺寸
stride	4	每层下采样步长
padding	0	卷积填充
total_stride	16	总下采样倍数
activation	nn.GELU	激活函数类型
interpol_size	None	若给定，先双线性插值到该尺寸再编码

其内部结构是若干 Conv2d + LayerNorm2d + GELU 单元组成的 nn.Sequential。层数由 num_layers = int(math.log2(total_stride) // math.log2(stride)) 推导并断言 stride**num_layers == total_stride；每层把通道数放大 stride² 倍，最后再接一个 1×1 卷积将通道对齐到 embed_dim。文档示例表明，输入 (1,1,256,256) 的 mask 经 MaskDownSampler(embed_dim=256, kernel_size=4, stride=4, padding=0, total_stride=16) 处理后输出为 (1,256,16,16)。在 encoders.py 第 379 行 的 SAM 2 mask 编码器中，它以 kernel_size=3, stride=2, padding=1, interpol_size=interpol_size 的配置被实例化。

2. CXBlock：ConvNeXt 风格的高效特征提取块

CXBlock 是修改版 ConvNeXt block，以卷积操作完成高效特征提取（blocks.py 第 128-210 行）。前向过程为：

input
 → depthwise conv (7×7, groups=dim)
 → LayerNorm2d(eps=1e-6)
 → permute 为 (N,H,W,C)
 → pwconv1: Linear(dim → 4*dim)
 → GELU
 → pwconv2: Linear(4*dim → dim)
 → gamma 逐通道缩放（layer scale，可学习）
 → 残差相加 input + DropPath(x)

参数要点：

• dim：通道数；kernel_size=7、padding=3 为大感受野 depthwise 卷积；
• use_dwconv：为 True 时按 groups=dim 做真正的 depthwise conv；
• layer_scale_init_value=1e-6：控制可学习缩放参数 gamma 的初值，为 0 时禁用该分支；
• 1×1 卷积以 nn.Linear 实现并配合通道转置，兼顾表达力与实现效率。

在 encoders.py 第 382 行，CXBlock(dim=256) 被作为 Fuser 的基底层用于融合图像特征与 mask 特征。

3. Fuser：多层同构特征融合器

Fuser 通过 copy.deepcopy 复制 num_layers 份相同 layer 依序处理张量，并可选地先经过一个 1×1 卷积输入投影（blocks.py 第 213-258 行）：

fuser = Fuser(layer, num_layers, dim=None, input_projection=False)

• layer：待复制的基底模块（通常是 CXBlock）；
• num_layers：堆叠层数；
• dim + input_projection=True：启用 nn.Conv2d(dim, dim, 1) 输入投影。

典型用法为 encoders.py 第 382 行 的 Fuser(CXBlock(dim=256), num_layers=2)，即用两层 ConvNeXt 块融合下采样后的 mask token 与像素特征。

三、双路注意力 Transformer：SAM2TwoWayAttentionBlock 与 SAM2TwoWayTransformer

1. SAM2TwoWayAttentionBlock：token 与 image 的双向往返注意力

SAM2TwoWayAttentionBlock 继承自 transformer.py 中的 TwoWayAttentionBlock，是 SAM/SAM 2 提示解码器的核心运算单元。它依次执行四步（对应基类 forward，transformer.py 第 194-237 行）：

1. sparse 自注意力：对 prompt token（queries）做 self-attention，随后 LayerNorm；
2. token→image 交叉注意力：queries 作为 q、image embedding 作为 k/v，让稀疏提示读取稠密图像特征；
3. MLP 块：对 queries 做非线性变换；
4. image→token 交叉注意力：交换 q/k 角色，让稠密图像特征反向吸收 token 信息（即"two-way"双向的含义）。

构造参数包括 embedding_dim（嵌入维度）、num_heads（头数，须能整除嵌入维度）、mlp_dim（MLP 隐层，默认 2048）、activation（默认 ReLU）、attention_downsample_rate（注意力内部维度降采样率，默认 2）、skip_first_layer_pe（首层是否跳过位置编码）。

该子类相对基类的差异在 blocks.py 第 313 行：将基类使用的 MLPBlock 替换为双层的 MLP(embedding_dim, mlp_dim, embedding_dim, num_layers=2, act=activation)，为 SAM 2 提供更强的 token 特征变换能力。

2. SAM2TwoWayTransformer：完整的解码 Transformer

SAM2TwoWayTransformer 在 transformer.py 基类 TwoWayTransformer 之上，将深度为 depth 的堆叠块全部替换为 SAM2TwoWayAttentionBlock，并约定首层 skip_first_layer_pe=True（blocks.py 第 316-378 行）。

其整体前向流程（基类 forward，transformer.py 第 81-122 行）：

image_embedding (B,C,H,W) → flatten → (B, H*W, C)      # keys
point_embedding (B,N,C)                                 # queries
重复 depth 次: SAM2TwoWayAttentionBlock(queries, keys, query_pe, key_pe)
最后: final_attn_token_to_image + LayerNorm 输出 queries
返回 (queries, keys)

transformer.py 中的底层 Attention 类（transformer.py 第 240-344 行）实现了带 downsample_rate 内部降维的多头注意力：当 downsample_rate > 1 时，internal_dim = embedding_dim // downsample_rate，可显著压缩交叉注意力计算量；同时通过 kv_in_dim 支持 q/k、v 维度不一致的场景。

SAM2TwoWayTransformer(depth=5, embedding_dim=256, num_heads=8, mlp_dim=2048) 正是 SAM 2 提示解码器实际使用的典型配置，最终输出 (1, N, 256) 的 token 特征与 (1, 256, H, W) 的图像特征，前者驱动 mask 与 score 预测。

四、RoPE 旋转位置编码注意力：RoPEAttention

RoPEAttention 继承自 Attention，将 Rotary Position Encoding 融入注意力机制（blocks.py 第 381-454 行），主要用于 SAM 2 的视频记忆注意力模块 memory_attention.py。

构造参数与机制：

• rope_theta=10000.0：频率基（theta），控制旋转频率衰减；
• feat_sizes=(32,32)：预计算频率张量的网格尺寸（对应 512 分辨率下 stride 16 的特征）；
• rope_k_repeat=False：交叉注意力中若 key 序列比 query 长，是否需要沿序列维重复 query 的旋转频率以对齐 k 长度。

forward 内部通过 partial(compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta) 计算轴向复指数频率，再调用 utils.py 中的 apply_rotary_enc 在复数域对 q/k 注入相对位置信息，随后交给 PyTorch 的 F.scaled_dot_product_attention 完成高效注意力。同时，freqs_cis 会依据当前 q 的 token 数自适应重算（if self.freqs_cis.shape[0] != q.shape[-2]），保证不同分辨率输入下编码依然有效；num_k_exclude_rope 参数则用于对部分 k token（如无位置意义的记忆 token）跳过旋转编码。

五、多尺度注意力：do_pool、MultiScaleAttention、MultiScaleBlock

这一组模块支撑了 Hiera（SAM 2 的高效分层图像编码器）的多尺度视觉 Transformer 骨干。

1. do_pool：带通道置换的池化工具

do_pool(x, pool, norm=None)（blocks.py 第 457-469 行）处理视觉 Transformer 中常见的 (B,H,W,C) 内存布局：先将张量置换为 (B,C,H,W) 执行池化，再置换回 (B,H',W',C)，可选地接一个 norm。它被 MultiScaleAttention（q_pool）与 MultiScaleBlock（shortcut 投影）共同使用。

2. MultiScaleAttention：带 Query 池化的多尺度自注意力

MultiScaleAttention（blocks.py 第 472-548 行）实现多尺度自注意力：

• qkv = nn.Linear(dim, dim_out * 3) 一次性投影出 q/k/v，随后 reshape(B, H*W, 3, nHead, -1) 拆头；
• 若配置 q_pool（即阶段切换处的下采样池化），则只对 q 做 do_pool 降采样，k/v 保持原分辨率——这是 Hiera 在 stage 边界压缩序列长度的关键设计；
• 注意力得分缩放为 scale = head_dim**-0.5，最终经 F.scaled_dot_product_attention 计算并由 proj 投影输出。

3. MultiScaleBlock：窗口化 + 阶段切换的完整 Block

MultiScaleBlock（blocks.py 第 551-662 行）将窗口注意力、query 池化、MLP 与残差组织为完整单元：

• window_size > 0 时，调用 utils.py 的 window_partition 将特征切分为不重叠窗口，注意力结束后再以 window_unpartition 还原（pad 信息通过 pad_hw 传递）；
• q_stride 非空时构造 nn.MaxPool2d 作为 stage 之间的空间降采样；窗口尺寸会同步按 window_size // q_stride[0] 折算；
• dim != dim_out 时通过 self.proj = nn.Linear(dim, dim_out) 并配合 do_pool 对齐 shortcut 分支；
• MLP 采用 MLP(dim_out, int(dim_out * mlp_ratio), dim_out, num_layers=2, act=act_layer)（即两层感知机）。

该块与 encoders.py 第 735 行 处 Hiera 编码器的 MultiScaleBlock 实例化一一对应，是 SAM 2 Hiera 骨干自上而下多尺度建模的主力构件。

六、位置编码：PositionEmbeddingSine 与 PositionEmbeddingRandom

1. PositionEmbeddingSine：正弦位置编码（面向网格）

PositionEmbeddingSine（blocks.py 第 665-776 行）为 2D 特征图生成经典的正弦/余弦位置编码：

• num_pos_feats：特征数，内部取半并断言为偶数（最终通道数 = 2 × num_pos_feats / 2 × 2 维交织）；
• temperature=10000：控制正弦频率随维度的衰减速度；
• normalize=True + scale：归一化到 [0, 2π]（scale 缺省为 2π），使编码与具体分辨率无关；若传 scale 而未开 normalize 会直接抛出 ValueError；
• 编码以 cache 字典按 (H,W) 缓存，相同分辨率只计算一次。

此外该类暴露了三个实用 API：

• encode_boxes(x, y, w, h)：把包围盒坐标与尺寸编码成位置向量（encode 为其向后兼容别名）；
• encode_points(x, y, labels)：对一组 2D 点做编码并拼接标签；
• forward(x)：为整张特征图生成 (B, 2*num_pos_feats, H, W) 的编码。

2. PositionEmbeddingRandom：随机空间频率编码（面向坐标）

PositionEmbeddingRandom（blocks.py 第 779-842 行）采用 Fourier 特征思想，用一个随机高斯矩阵把坐标映射到高频正弦/余弦空间：

• 构造时以 scale * torch.randn((2, num_pos_feats)) 注册 positional_encoding_gaussian_matrix 缓冲区（scale <= 0 时回退为 1.0）；
• _pe_encoding 将 [0,1]² 坐标线性变换到 [-1,1]² 后与高斯矩阵相乘、乘 2π，输出 [sin, cos] 拼接；
• forward(size) 对 (h,w) 网格做累积坐标生成（cumsum - 0.5 再除以 h/w），产出 (C,H,W) 编码；
• forward_with_coords(coords, image_size) 供提示编码器对任意坐标（点/框）即时编码，SAM 中点提示与框提示的编码正是此路径。

同时该类在初始化时显式关闭确定性算法（torch.use_deterministic_algorithms(False)、cudnn.deterministic = False），注释表明是为规避 cumsum_cuda_kernel 的不确定性报错。

七、通用 ViT 构建块：Block、REAttention 与 PatchEmbed

最后一组模块是通用视觉 Transformer 的三大件，也是 SAM 1 / MobileSAM 骨干（TinyViT）的基础零件。

1. Block：支持窗口注意力的标准 Transformer 块

Block（blocks.py 第 845-933 行）结构为 LayerNorm → (窗口化) REAttention → 残差 → MLPBlock。参数 window_size 决定注意力模式：为 0 时全局注意力，大于 0 时先 window_partition 再在窗口内做注意力、随后 window_unpartition 还原。use_rel_pos=True 时启用相对位置编码，此时 input_size 决定可学习的 rel_pos_h/rel_pos_w 参数长度。

该块在 encoders.py 第 103-114 行 被 ImageEncoderViT 依深度循环实例化，且支持 global_attn_indexes——被索引命中的层强制全局注意力（window_size=0），其余层走窗口注意力，这是 SAM 1 ViT 编码器的分层混合策略。

2. REAttention：相对位置注意力

REAttention（blocks.py 第 936-1015 行）是多头自注意力加可选相对位置编码的实现：

• qkv 一次投影 (dim → 3*dim)，重塑为 (3, B, nHead, H*W, C) 后拆分；
• 注意力得分 attn = (q * scale) @ k^T；
• use_rel_pos=True 时调用 utils.py 的 add_decomposed_rel_pos，按 MVITv2 的解耦式分解把 height/width 两轴相对位置偏置（rel_pos_h、rel_pos_w，长度 2*input_size-1）注入得分矩阵；两参数以 rel_pos_zero_init=True 默认零初始化，保证训练早期稳定性；
• 注意力经 softmax 后作用于 v，proj 输出。因为 Block 的 REAttention 使用的是直接矩阵注意力而非 SDPA，与 RoPEAttention/MultiScaleAttention 形成实现上的对照。

3. PatchEmbed：图像到 patch 嵌入

PatchEmbed（blocks.py 第 1018-1067 行）是视觉 Transformer 的第一层：用单个 nn.Conv2d(in_chans, embed_dim, kernel_size, stride, padding, bias) 将图像切分为 patch 并投影为嵌入，forward 再把结果从 (B,C,H,W) 转置为 (B,H,W,C) 以匹配注意力/MLP 的内存布局。默认 kernel_size=(16,16), stride=(16,16), embed_dim=768 时，224×224 输入产出 14×14 的 patch 序列。它同时被 SAM 1 的 ImageEncoderViT 与 SAM 3 的 vitdet.py 复用。

八、工程衔接：这些模块如何进入实际模型

要观察这些基础构件在真实模型中的拼装，可以从两条链路入手：

• SAM / MobileSAM 链路：build.py 依据配置选择 vit_h/vit_b/... 或 TinyViT（tiny_encoder.py）构造图像编码器，配套 encoders.py 的 ImageEncoderViT，解码端由 sam.py 引用 SAM2TwoWayTransformer（SAM 2 使用）与 transformer.py 的 TwoWayTransformer（SAM 1 使用）；
• SAM 2 视频链路：memory_attention.py 通过 RoPEAttention 在帧间记忆 token 上做带旋转位置编码的注意力，保证多帧时序位置的一致性。

若要在工程中复用作独立层，可直接仿照模块 docstring 中的最小示例：

import torch
from ultralytics.models.sam.modules.blocks import (
    Block, CXBlock, DropPath, MaskDownSampler, MultiScaleBlock,
    PositionEmbeddingRandom, PositionEmbeddingSine,
)

# 训练期以 0.2 概率随机丢弃残差分支
drop_path = DropPath(drop_prob=0.2, scale_by_keep=True)
out = drop_path(torch.randn(32, 64, 224, 224))

# ConvNeXt 特征提取（通道数不变的残差块）
cx = CXBlock(dim=64, kernel_size=7, padding=3)
out = cx(torch.randn(1, 64, 56, 56))

# mask 编码：从单通道 mask 得到 embed_dim 特征
down = MaskDownSampler(embed_dim=256, kernel_size=4, stride=4, padding=0, total_stride=16)
out = down(torch.randn(1, 1, 256, 256))

# 位置编码
pe = PositionEmbeddingSine(num_pos_feats=128)
emb = pe(torch.randn(1, 3, 224, 224))            # (1, 128, 224, 224)
pe_r = PositionEmbeddingRandom(num_pos_feats=64)
enc = pe_r((32, 32))                              # (128, 32, 32)

结语

从正则化工具 DropPath，到卷积风格的 MaskDownSampler/CXBlock/Fuser，再到双路解码器 SAM2TwoWayTransformer、带旋转编码的 RoPEAttention、多尺度注意力 MultiScaleAttention/MultiScaleBlock，以及两类位置编码和通用 ViT 构件 Block/REAttention/PatchEmbed，ultralytics/models/sam/modules/blocks.py 以极简而自洽的模块划分，支撑起 SAM、SAM 2、SAM 3 及轻量 MobileSAM 多条模型线。阅读这些代码的每一处 docstring 示例与断言，可以直观看到其复用性设计意图——同一份 PatchEmbed 同时服务图像级 ViT 与 patch 化骨干，同一份 DropPath 在全部深层网络间共享。理解这些构件，是进一步深入 SAM 系列提示解码、视频记忆与高效分层编码机制的最佳起点。