Ultralytics SAM3 视觉 Transformer 编码器源码深度解析：`encoder.py` 中的注意力层与多模态融合

本文以 Ultralytics 仓库中 SAM3 编码器模块 的官方 API 参考文档（encoder.md）为主线，系统讲解 TransformerEncoderLayer、TransformerEncoder、TransformerEncoderFusion 与 pool_text_feat 四组核心符号的架构设计、逐行数据流与实际构造参数。读者将掌握 SAM3 视觉-语言 Transformer 编码器如何在自注意力与交叉注意力两级交互中把文本 Prompt 注入图像特征，为后续解码器产出高质量 memory。

一、模块定位：SAM3 中的"图像编码 + 文本注入"枢纽

SAM3 是 Ultralytics 推出的支持视觉语言联合建模的分割模型，其整体前向流程大致为：视觉主干（多分辨率特征）→ Neck → Transformer 编码器 → Transformer 解码器 → 分割头。本参考文档对应源码文件 ultralytics/models/sam/sam3/encoder.py，其中定义的类承载了"编码器"这一中间枢纽的职责：

• 它接收来自主干网络的多层级视觉特征，将其展平为 token 序列；
• 它同时接收文本/几何 Prompt（如 CLIP 文本 token 或点击坐标），并让图像 token 通过交叉注意力去"读取"这些 prompt；
• 它最终把编码后的图像 memory 与文本特征一并交给解码器做 grounding 与掩码预测。

从调用关系看，该文件导出的 TransformerEncoderFusion、TransformerEncoderLayer 在 build_sam3.py 中被导入并用于 _create_sam3_transformer() 组装编码器实例；而 sam3_image.py 中 SAM3SemanticModel._run_encoder() 则通过 self.transformer.encoder(...) 实际驱动前向过程。因此，把这份参考文档对应的源码读懂，就等于理解了 SAM3 中"图像特征如何与文本语义互相校准"的机制。

参考文档本身通过 ::: ultralytics.models.sam.sam3.encoder.xxx 语法列出的四个公开符号及其 docstring 即构成了 API 全貌，下文逐一展开。

二、TransformerEncoderLayer：自注意力 + 图像交叉注意力的最小单元

参考文档列出的第一个类 TransformerEncoderLayer 是整个编码器的"原子构建块"。如其 docstring 所述，该类"performs self-attention followed by cross-attention"，即对输入序列先做自注意力、再做与另一路输入（通常是图像特征）的交叉注意力。该层最初曾命名为 TransformerDecoderLayer，后因其更准确的架构角色而被重命名。

2.1 构造参数：三个位置编码开关与归一化策略

该类构造签名如下（见 encoder.py 的 docstring）：

参数	含义	实际取值（来自 build_sam3.py）
d_model	模型隐藏维度	256
dim_feedforward	FFN 中间层维度	2048
dropout	Dropout 概率	0.1
pos_enc_at_attn	自注意力 Q/K 是否叠加位置编码	True
pos_enc_at_cross_attn_keys	交叉注意力 Key 是否叠加 memory 的位置编码	False
pos_enc_at_cross_attn_queries	交叉注意力 Query 是否叠加位置编码	False
pre_norm	使用 Pre-Norm（True）还是 Post-Norm（False）	True
self_attention / cross_attention	自定义注意力模块，缺省时创建默认多头注意力	8 头、embed_dim=256

在 build_sam3.py 的真实构造中，self_attention 与 cross_attention 均显式传入 nn.MultiheadAttention(num_heads=8, dropout=0.1, embed_dim=256, batch_first=True)，未依赖默认分支。层内部包含三条残差支路（自注意力、交叉注意力、前馈网络），每条支路配有独立的 LayerNorm（norm1/norm2/norm3）与 Dropout（dropout1/dropout2/dropout3），FFN 采用 Linear → ReLU → Dropout → Linear 结构，激活函数固定为 nn.ReLU()。

此外层上还维护一个 layer_idx 属性（初始为 None），用于在深层编码器中让某些层依据自身处于第几层来决定行为（例如只在特定层对记忆库做交叉注意力）。

2.2 前向实现：Post-Norm 与 Pre-Norm 两条路径

该层支持两种归一化顺序，forward() 依据 self.pre_norm 在 forward_pre 与 forward_post 间分发（见 encoder.py）。

Post-Norm 路径 forward_post（见 encoder.py）的核心计算流为：

q = k = tgt + query_pos if self.pos_enc_at_attn else tgt
# ① 自注意力
tgt2 = self.self_attn(q, k, value=tgt, ...)[0]
tgt = self.norm1(tgt + self.dropout1(tgt2))
# ② 对图像的交叉注意力
tgt2 = self.cross_attn_image(
    query=tgt + query_pos if self.pos_enc_at_cross_attn_queries else tgt,
    key=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
    value=memory, ...)[0]
tgt = self.norm2(tgt + self.dropout2(tgt2))
# ③ FFN
tgt = self.norm3(tgt + self.dropout3(self.linear2(self.dropout(self.activation(self.linear1(tgt))))))

其中 query_pos 一般承载着逐 token 的位置编码，pos 则为 memory（交叉注意力的 key/value 来源）携带位置信息。

Pre-Norm 路径 forward_pre（见 encoder.py）则先归一化再进注意力，且在结构上多了对 dac（Divide-and-Conquer）模式的特殊处理：当 dac=True 时，仅对 query 的前半部分做自注意力，后半部分原样保留，完成后再 torch.cat 拼接回去。这一技巧能显著削减自注意力的计算量，符合超长序列场景下的分治设计。同样注意 Pre-Norm 路径在自注意力前显式 .contiguous()，并把 memory 转成与归一化后 tgt2 一致的 dtype，规避低精度下类型不匹配问题。

由于 SAM3 编码器层在 build_sam3.py 中设置为 pre_norm=True，实际前向会走 forward_pre 分支。

三、TransformerEncoder：多层级特征展平与多层堆叠

TransformerEncoder 将上一节的最小单元堆叠起来，专门处理来自主干网络的多分辨率特征（见 encoder.py 的类 docstring）。

3.1 构造要点

TransformerEncoder(
    layer: nn.Module,        # 待克隆的编码器层
    num_layers: int,         # 堆叠层数
    d_model: int,            # 隐藏维度
    num_feature_levels: int, # 处理的特征层数
    frozen: bool = False,    # 是否冻结全部参数
    use_act_checkpoint: bool = False,  # 训练时是否启用激活检查点
)

• 多层克隆复用 ultralytics/nn/modules/utils.py 中的 _get_clones（nn.ModuleList([copy.deepcopy(module) for _ in range(n)])），各层参数相互独立；
• 当 num_feature_levels > 1 时创建可学习的分层嵌入 level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model))，用于区分来自不同分辨率的 token；
• frozen=True 时将所有 requires_grad 置为 False，便于把编码器当作冻结的预训练组件；
• 构造末尾会为每一层写入 layer_idx = enumerate(layers) 的下标；
• 前向中若处于 self.training，会断言 self.use_act_checkpoint 必须开启（源码注释为 "activation ckpt not enabled in encoder"），用于约束大模型训练的内存占用。

3.2 _prepare_multilevel_features：把特征图变成 token 序列

_prepare_multilevel_features（见 encoder.py）负责将多个 (bs, c, h, w) 特征图统一扁平化：

1. 断言特征层数等于 num_feature_levels；
2. 对每层 src、mask、pos_embed 分别做 flatten(2).transpose(1, 2)，得到 (bs, hw, c) 的 token 序列；
3. 若存在分层嵌入，则 lvl_pos_embed = pos_embed + level_embed[lvl]；
4. 将各层结果沿序列维 torch.cat 拼成 (bs, Σhw, c)；
5. 生成 spatial_shapes（各层高宽）、level_start_index（各层 token 起始下标），并借助 model_misc.py 的 get_valid_ratio(mask) 计算每层高宽有效比例 valid_ratios，供解码器的可变形/稀疏注意力定位参考点使用。

3.3 forward：把 Prompt 当作交叉注意力的 memory

forward（见 encoder.py）的签名在普通图像编码器之上加入了 prompt 通道：

def forward(self, src, src_key_padding_masks=None, pos=None,
            prompt=None, prompt_key_padding_mask=None, encoder_extra_kwargs=None):

循环堆叠各层时，每层收到的 layer_kwargs 为：memory=prompt（交叉注意力的 key/value）、memory_key_padding_mask=prompt_key_padding_mask、query_pos=lvl_pos_embed_flatten（即图像侧展平位置编码）、tgt=output。也就是说：这里的交叉注意力让每个图像 token 去查询文本/几何 prompt token，从而把语义信息逐层写进视觉特征。encoder_extra_kwargs 支持向每一层透传额外参数（可复用同一份配置）。

返回时统一转换为 seq-first 布局并打包为六元组：(output, key_padding_masks_flatten, lvl_pos_embed_flatten, level_start_index, spatial_shapes, valid_ratios)。

四、TransformerEncoderFusion：文本与图像特征的真正融合点

TransformerEncoderFusion 继承自 TransformerEncoder，参考文档明确其职责是"fuses text and image features"。它额外引入一条显式的跨模态融合捷径：把池化后的文本特征直接加到图像特征上（见 encoder.py）。

4.1 新增构造参数

参数	含义	SAM3 主模型实际取值
add_pooled_text_to_img_feat	是否把池化文本加到图像特征	False
pool_text_with_mask	池化文本时是否依据 padding mask	True
compile_mode	torch.compile 编译模式，None 表示关闭	由外层 compile 开关决定

当 add_pooled_text_to_img_feat=True 时，会额外创建投影层 self.text_pooling_proj = nn.Linear(d_model, d_model)，用于对池化后的文本向量做一次线性变换再广播到图像侧。若传入非空 compile_mode，则用 self.forward = torch.compile(self.forward, mode=compile_mode, fullgraph=True) 对整个前向做图编译加速（见 encoder.py）。

值得注意的是 build_sam3.py 为 SAM3 图像模型实例化 TransformerEncoderFusion 时的默认配置是：6 层、d_model=256、num_feature_levels=1、frozen=False、use_act_checkpoint=True、add_pooled_text_to_img_feat=False、pool_text_with_mask=True。这从源码层面印证了：SAM3 主干默认并不走"文本全局池化加到图像"这条简单捷径，而是依赖 2.2 与 3.3 中逐层交叉注意力完成语义注入，融合类提供的捷径仅在需要时以开关形式启用。

4.2 前向差异与输出字典

TransformerEncoderFusion.forward 相对父类的主要差异：

1. 恢复视觉空间形状：传入的 src/src_pos/mask 可能是 seq-first 的展平结果，配合 feat_sizes 列表可执行 reshape(h, w, bs, -1).permute(2, 3, 0, 1) 还原为 (bs, c, h, w)（见 encoder.py），随后交给父类统一处理；
2. 显式文本-图像融合：若启用，则 pooled_text = pool_text_feat(prompt, prompt_key_padding_mask, self.pool_text_with_mask) 后经 text_pooling_proj 投影为 (bs, d_model, 1, 1)，再 src = [x.add_(...) for x in src] 原地加到每个图像特征图上（见 encoder.py），等价于给所有图像 token 注入统一的"全局语义偏置"；
3. 输出结构化字典：返回 memory（编码后图像特征）、padding_mask、pos_embed、memory_text（透传 prompt，即编码后的文本侧特征）、level_start_index、spatial_shapes、valid_ratios 七项。memory_text 正是后续解码器做文本交叉注意力所需的"文本 memory"。

4.3 在 SAM3SemanticModel 中的驱动方式

在 sam3_image.py 的 _run_encoder 中，调用形态为：

memory = self.transformer.encoder(
    src=img_feats.copy(),      # 复制列表，避免编码器 in-place 修改
    src_key_padding_mask=None,
    src_pos=img_pos_embeds.copy(),
    prompt=prompt,             # 文本 + 几何 prompt（seq-first）
    prompt_key_padding_mask=prompt_mask,
    feat_sizes=vis_feat_sizes,
    encoder_extra_kwargs=encoder_extra_kwargs,
)

返回的 memory["memory"] 作为 encoder_hidden_states 喂给解码器，memory["memory_text"]（若存在）作为编码后的 prompt 一并传递，从而完成"编码器输出 = 图像 memory + 文本 memory"的解耦表达。

五、pool_text_feat：带掩码的均值池化函数

模块级函数 pool_text_feat(prompt, prompt_mask, pool_with_mask) 是 Fusion 编码器与池化文本捷径共用的工具（见 encoder.py），其 docstring 为 "Mean-pool the prompt embeddings over the valid tokens only"。

实现要点：

• 输入 prompt 形状为 (seq, bs, dim)（seq-first），prompt_mask 形状为 (bs, seq)，其中 False 表示有效 token、True 表示 padding；
• pool_with_mask=False 时直接返回 prompt.mean(dim=0)，即对全部 token（含 padding）做无差别均值；
• pool_with_mask=True 时先翻转掩码得到 is_valid = (~prompt_mask).float().permute(1, 0)[..., None]，统计 num_valid = clamp(sum(is_valid), min=1.0) 防止除零，再 pooled_text = (prompt * is_valid).sum(dim=0) / num_valid，即仅在有效 token 上求平均。

有趣的是，model_misc.py 的 DotProductScoring.mean_pool_text 实现了完全一致的掩码均值池化逻辑，说明"按有效 token 池化 prompt"是 SAM3 中用于打分与特征融合的通用设计，二者互为印证。

六、从配置到前向的完整调用链（自查清单）

为便于读者基于当前仓库快速验证，以下为阅读顺序建议：

1. 官方 API 参考：docs/en/reference/models/sam/sam3/encoder.md，确认四个导出符号；
2. 组件实现：ultralytics/models/sam/sam3/encoder.py，重点看 Pre-Norm 分支与 Fusion 的原地相加逻辑；
3. 编码器组装：ultralytics/models/sam/build_sam3.py#L72-L132 中 _create_sam3_transformer()，核对 d_model=256 / dim_feedforward=2048 / num_layers=6 等生产参数；
4. 高层驱动：ultralytics/models/sam/sam3/sam3_image.py#L130-L164 的 _run_encoder，观察 prompt 拼接与字典消费；
5. 支撑工具：_get_clones（ultralytics/nn/modules/utils.py#L15-L32）与 get_valid_ratio（ultralytics/models/sam/sam3/model_misc.py#L130-L138）。

需要强调的是，该编码器设计脱胎于 GroundingDINO（源码头部注明 "Based on https://github.com/IDEA-Research/GroundingDINO"，版权归属 Meta Platforms），因此其"自注意力校准 + 图像-文本交叉注意力 + 可选池化捷径"的三段式结构，本质上是检测/分割中 grounding 范式向 SAM 系列交互式分割演进的产物。理解 encoder.py，也就把握住了 SAM3 中视觉-语言融合的技术内核。

七、小结

回到参考文档本身，encoder.md 通过四个 ::: 引用块精炼地划定了本模块的公开 API 边界。将其与仓库源码对照后可以看到：TransformerEncoderLayer 是"自注意力 + 图像交叉注意力 + FFN"的归一化可配置最小单元；TransformerEncoder 负责多层级特征展平与多层堆叠、并为每层注册 layer_idx；TransformerEncoderFusion 在父类基础上增加了"池化文本注入图像"的捷径与 torch.compile 支持，并以字典形式把图像 memory 与文本 memory 同时交给解码器；pool_text_feat 则保证了池化过程对 padding token 的鲁棒性。这些机制共同保证了 SAM3 在视觉-语言联合推理中既能保留图像的细粒度空间信息，又能稳定地吸收文本语义约束。