Ultralytics SAM3 几何编码器源码剖析：从几何 Prompt 到 Transformer 可读序列的完整实现解读

Ultralytics SAM3 几何编码器源码剖析从几何 Prompt 到 Transformer 可读序列的完整实现解读【免费下载链接】ultralyticsUltralytics YOLO26, YOLO11, YOLOv8 — object detection, instance segmentation, semantic segmentation, image classification, pose estimation, object tracking项目地址: https://gitcode.com/GitHub_Trending/ul/ultralytics文章导读本文基于 Ultralytics 仓库 geometry_encoders.py 的公开接口与真实实现深入剖析 SAM3Segment Anything Model 3中几何提示geometric prompt的表示与编码链路Prompt数据结构如何承载 box/point 几何提示is_right_padded与concat_padded_sequences两个工具函数如何维护右填充right-padded的不定长批序列以及SequenceGeometryEncoder如何将归一化的 CxCyWH 框投影、RoI 池化并融合位置编码后送入 Transformer。读完本文你将能够准确理解 geometry_encoders.md 所列四个公开符号的输入输出约定、维度规则与参数语义并能在二次开发接入自定义几何提示、调整编码配置时直接定位到相关源码与调用点。说明docs/en/reference/...下以::: 模块路径.符号形式呈现的是由 mkdocstrings 从源码 docstring 自动生成的 API 参考页其技术细节的权威来源正是本文所引用的源码模块与上游调用点。一、模块定位SAM3 中框提示通往 Transformer 的必经之路ultralytics/models/sam/sam3/geometry_encoders.py位于 SAM3 子包内是标准 SAM3 语义模型SAM3SemanticModel中文本提示 几何提示双通道提示编码的关键一环。从 sam3_image.py 可以看到模型在编码阶段先将几何提示交给self.geometry_encoder即SequenceGeometryEncoder实例得到geo_feats, geo_masks再与视觉提示 embedding 拼接成完整的 prompt 序列随后被 Transformer encoder 消费# ultralytics/models/sam/sam3/sam3_image.py geo_feats, geo_masks self.geometry_encoder( geo_promptgeometric_prompt, img_featsimg_feats, img_sizesvis_feat_sizes, img_pos_embedsimg_pos_embeds, ) prompt torch.cat([geo_feats, visual_prompt_embed], dim0) prompt_mask torch.cat([geo_masks, visual_prompt_mask], dim1)由此可见该模块在整个 SAM3 前向中的位置图像特征由 backbone 产出几何提示框/角点由该模块编码为与图像特征同维的 token 序列供 Transformer 在 cross-attention 中作为 query/condition 使用。在调用链另一侧build_sam3.py 以模块配置的方式给出了标准 SAM3 语义模型的几何编码器组装实例详见第四节predict.py 在推理器SAM3Predictor中通过Prompt(box_embeddingstorch.zeros(0, num_prompts, 4), ...)构造零框 dummy 几何提示用于无框提示时的推理路径。这两处共同构成该模块在模型构建与推理两条链路上的真实用法证据。二、Prompt统一承载几何提示的数据容器Prompt在 geometry_encoders.py 的类 docstring 中被定位为操作几何提示的工具类。它是介于上游归一化坐标输入与下游编码器之间的统一数据容器其维度约定如下源码注释原文要点张量形状含义允许取值box_embeddingsN_boxes × B × C_box每个框的几何特征归一化框坐标或预计算 embeddingbox_maskB × N_boxesattention maskPyTorch 约定1 表示被 mask/padNone表示无 mask 项box_labelsN_boxes × B正/负样本标签long 型None表示全部视为正样本三个关键约定在类 docstring 中被强调序列维度在前所有序列张量按 PyTorch 约定组织——序列长度N在前、batch 维度B在后而 mask 张量则 batch-first。盒坐标采用归一化 CxCyWH 格式即(center_x, center_y, width, height)坐标值归一化到[0, 1]这一点由SequenceGeometryEncoder的类 docstring 声明并由构造函数断言box_embeddings.shape[-1] 4来保证见 geometry_encoders.py。标签默认全正、mask 默认全非 padbox_labels缺省时填充torch.ones(...)box_mask缺省时填充torch.zeros(...)geometry_encoders.py。2.1 构造与形状校验Prompt.__init__接受box_embeddings / box_mask / box_labels三个可选参数。当box_embeddings is None时构造出一个空 prompt三字段全为None可用于零框推理否则依次补齐缺失的labels与mask并对以下条件做断言校验box_embeddings前两维为[N_boxes, B]末维必须为 4四个几何量box_mask形状恰为[B, N_boxes]box_labels形状恰为[N_boxes, B]三者 device 一致。源码中这些校验一方面保证数据自洽另一方面与concat_padded_sequences的断言共同构成编码前最后一道维度防火墙。2.2 append_boxes动态追加框提示Prompt.append_boxes(boxes, labelsNone, maskNone)geometry_encoders.py支持两种场景首框初始化当self.box_embeddings is None时直接以本次boxes初始化并补齐 labels/mask追加既有校验 batch 大小与 shapes 一致后借助concat_padded_sequences分别对box_labels先unsqueeze(-1)再拼接、再squeeze(-1)与box_embeddings完成右填充序列拼接。该方法与第一节所述_get_dummy_prompt零框形成互补推理中通过逐帧追加框即可在保持右填充不变式的前提下累积提示。从源码结构看这类逐次追加的能力为视频/交互式场景下提示的增量维护提供了原语。三、两个基础工具函数右填充判定与不等长序列拼接3.1 is_right_padded判定 padding 是否在右侧def is_right_padded(mask: torch.Tensor): return (mask.long() torch.sort(mask.long(), dim-1)[0]).all()按 PyTorch 约定padding mask 中 1 表示被 pad 的占位。若 padding 位于序列右侧则 mask 中先出现一段 0、后出现一段 1恰好是**非递减升序**序列。该函数通过比较 mask 与其升序排序结果是否逐元素相等来判断整批序列是否都是右填充geometry_encoders.py。它被concat_padded_sequences内部以torch._assert调用作为拼接操作的前置不变量检查。3.2 concat_padded_sequences两条右填充序列的无缝拼接concat_padded_sequences(seq1, mask1, seq2, mask2, return_indexFalse)geometry_encoders.py是Prompt.append_boxes与SequenceGeometryEncoder.forward追加 CLS token 时共同依赖的核心拼接原语。其输入输出约定如下参数形状说明seq1(L1, B, H)序列优先、特征在末维mask1(B, L1)1 表示 padseq2(L2, B, H)同上mask2(B, L2)同上return_indexbool是否额外返回 seq2 在拼接序列中的索引实现要点算法层面前置断言核对 batch、hidden、序列长度两两匹配并断言mask1/mask2均为右填充。计算真实长度actual_seqN_lengths (~maskN).sum(dim-1)统计每样本非 pad 的真实 token 数拼接后每样本真实长度相加为final_lengths最大可能长度为max_length L1 L2。构造拼接 mask利用广播比较torch.arange(max_length) final_lengths生成新的右填充 mask——凡超过该样本真实总长度的位置都置 1。移位放置 seq2先新建(max_length, B, H)的全零张量把seq1直接放进前L1行随后计算 seq2 各行应落入的目标行号index arange(L2)[:,None] actual_seq1_lengths[None]即在 seq1 实际长度基础上偏移用scatter将 seq2 写入对应位置。可选返回 indexreturn_indexTrue时额外返回index形状(L2, B)用于从拼接序列中精确取回 seq2 的元素。正是mask 右填充 每样本真实长度可推导这一不变式使得该函数无需逐样本循环即可高效完成变长序列的批式拼接拼接结果天然仍是右填充序列可直接馈给下游 Transformer 的key_padding_mask。四、SequenceGeometryEncoder构造参数与三种框编码路径SequenceGeometryEncoder的完整 docstring 与构造函数位于 geometry_encoders.py。它声明接受归一化 CxCyWH格式的框框可被三种方式编码三者互不排斥、可叠加求和direct projection线性投影对 4 维坐标做线性投影到d_modelpoolingRoI align从 backbone 特征图做 RoI align汇聚框内区域特征pos encoder位置编码对框中心做正余弦位置编码复用PositionEmbeddingSine。作为替代方案框还可以被拆解为左上/右下两个角点来编码encode_boxes_as_pointsTrue。4.1 构造参数语义参数类型语义与影响encode_boxes_as_pointsbool是否把框拆为两个角点编码。True 时使用(左上, 右下)两组点boxes_direct_projectbool线性投影路径对应nn.Linear(4, d_model)boxes_poolboolRoI 路径对应nn.Conv2d(d_model, d_model, roi_size)boxes_pos_encbool位置编码路径对应nn.Linear(d_model 2, d_model)d_modelint模型宽度所有编码输出的公共通道维度pos_encnn.Module位置编码器如PositionEmbeddingSine用于框中心编码num_layersint后续 Transformer 编码层数量 0时强制建议开启 CLSlayernn.Module单个 Transformer 编码层由_get_clones深拷贝复制roi_sizeint7RoI align 输出尺寸高/宽add_clsboolTrue是否在序列头部加入可学习的 CLS tokenadd_post_encode_projboolTrue是否追加Linear LayerNorm作为编码后精化use_act_ckptboolFalse是否在多层编码器上启用激活检查点省显存构造函数中还蕴含两个与配置一致性相关的细节标签 embedding 数量动态化编码为框时每 token 只有正/负 2 类标签编码为角点时每点可能出现普通正负、左上正负、右下正负共 6 类故label_embed nn.Embedding(num_labels, d_model)中num_labels 6 if encode_boxes_as_points else 2geometry_encoders.py。非角点模式至少需要一种框编码方式若encode_boxes_as_pointsFalse且三种框编码开关全为 False则直接断言报错Error: need at least one way to encode boxes。RoI 相关模块附带输入归一化当任一 pooling 路径启用时img_pre_norm从nn.Identity()切换为nn.LayerNorm(d_model)在池化前对特征做逐层归一。4.2 仓库中的真实组装示例build_sam3.py 中build_semantic_sam3给出的标准配置可作为理解各参数的权威样例input_geometry_encoder SequenceGeometryEncoder( pos_encPositionEmbeddingSine( num_pos_feats256, normalizeTrue, scaleNone, temperature10000, ), encode_boxes_as_pointsFalse, boxes_direct_projectTrue, boxes_poolTrue, boxes_pos_encTrue, d_model256, num_layers3, layerTransformerEncoderLayer( d_model256, dim_feedforward2048, dropout0.1, pos_enc_at_attnFalse, pre_normTrue, pos_enc_at_cross_attn_queriesFalse, pos_enc_at_cross_attn_keysTrue, ), use_act_ckptTrue, add_clsTrue, add_post_encode_projTrue, )可以看到生产级语义模型默认三条框编码路径全部开启并求和d_model256、3 层 Transformer 编码层num_layers3、激活检查点开启use_act_ckptTrue且使用正弦位置编码PositionEmbeddingSine(num_pos_feats256, temperature10000)。其中_get_clones(layer, num_layers)负责把同一个layer深拷贝出多层堆叠该工具函数定义在 nn/modules/utils.py是 Ultralytics 内部通用的模块克隆助手。五、三种框编码方式的底层实现5.1 直接线性投影direct projection在_encode_boxes中geometry_encoders.py若boxes_direct_project开启则将归一化 4 维坐标直接送入nn.Linear(4, d_model)proj self.boxes_direct_project(boxes.to(img_feats.dtype))这里boxes.to(img_feats.dtype)表明坐标会先被转换成与图像特征一致的精度如混合精度下为 fp16。5.2 RoI align 特征汇聚pooling若boxes_pool开启其流程geometry_encoders.py为取img_feats的H, W用xywh2xyxy把归一化 CxCyWH 框转为xyxy再按[W, H, W, H]缩放反归一化到像素坐标调用torchvision.ops.roi_align延迟导入以加快ultralytics包加载在特征图上采样得到(B*N, d_model, roi_size, roi_size)的 RoI 特征用nn.Conv2d(d_model, d_model, roi_size)将每个 RoI 汇聚成d_model维向量roi_size7时等价于 7×7 全局卷积池化view(bs, n_boxes, d_model).transpose(0, 1)还原为(N, B, d_model)的序列优先格式。注意xywh2xyxy来自 utils/ops.py是该仓库全局复用的坐标转换工具与 SAM/SAM2 其他模块保持一致。5.3 框中心位置编码pos encoder若boxes_pos_enc开启geometry_encoders.py会把框解绑为cx, cy, w, h四个标量组调用位置编码器的encode_boxesenc self.pos_enc.encode_boxes(cx.flatten(), cy.flatten(), w.flatten(), h.flatten()) proj self.boxes_pos_enc_project(enc.to(img_feats.dtype))boxes_pos_enc_project为nn.Linear(d_model 2, d_model)——多出的 2 维来自encode_boxes末尾直接拼接的(h, w)原始宽度/高度。PositionEmbeddingSine.encode_boxes的实现在 nn/.../blocks.py注该文件真实路径为ultralytics/models/sam/modules/blocks.py其做法是把中心点按y, x顺序拼接正弦编码后再接h, wpos_x, pos_y self._encode_xy(x, y) return torch.cat((pos_y, pos_x, h[:, None], w[:, None]), dim1)三路结果直接投影 / RoI / 位置编码在_encode_boxes中通过先到先得、后到累加的方式求和最后统一加上type_embed self.label_embed(boxes_labels.long())的标签 embedding作为几何 token 的最终表示。这印证了类 docstring 中三种编码互不排斥、多选即求和的表述。六、角点编码模式encode_boxes_as_pointsTrue 的分支当encode_boxes_as_pointsTrue时forward走_encode_points分支geometry_encoders.py将每个框升级成一对角点 tokenboxes_xyxy xywh2xyxy(boxes)转为归一化 xyxy再split(split_size2, dim-1)拆成top_left前两维与bottom_right后两维对角点标签做偏移区分来源labels_tl boxes_labels 2、labels_br boxes_labels 4配合构造时预留的 6 类label_embed使 Transformer 能区分左上正/负与右下正/负两组点按序列维torch.cat拼接成(2*N, B, 2)的点序列mask 相应横向拼接交由_encode_pointsnn.Linear(2, d_model)直接投影 2 维坐标并叠加 6 类标签 embedding。这种模式下序列长度翻倍每框两 token换来的是模型对框两角位置更细粒度的关注适合需要更强空间定位能力的设定。七、forward 主流程CLS、归一化精化与 Transformer 编码层SequenceGeometryEncoder.forward(geo_prompt, img_feats, img_sizes, img_pos_embedsNone)geometry_encoders.py的完整流水线为取数从geo_prompt解出boxes / boxes_mask / boxes_labels同时取img_feats[-1]作为序列优先H*W, B, C的跨模态记忆供后续 cross-attention 使用。池化前的特征准备若启用了任一条 pooling 路径则用img_pre_normLayerNorm对最后一层图像特征归一化并由(H*W, B, C)重排为(N, C, H, W)图像格式以配合 RoI align。按模式编码encode_boxes_as_points为 True 走角点路径否则走框编码路径得到final_embeds (L, B, d_model)与final_mask (B, L)。追加 CLS若add_clsTrue用可学习cls_embed生成 1 个全 batch 共享的 CLS tokenmask 位为 0永不被 pad并通过concat_padded_sequences将其拼接在序列头部geometry_encoders.py。这也解释了构造器中使用 Transformer 时强烈建议开启 CLS的断言——CLS 是编码层输出汇聚的聚合位。后编码精化若final_proj存在执行norm(final_proj(final_embeds))Linear LayerNorm。堆叠 Transformer 编码层将num_layers个克隆的layer逐层作用——每层以图像特征为memory、几何序列为tgttgt_key_padding_mask传入右填充 mask、pos传入图像侧位置编码最终经encode_normLayerNorm输出。返回的(final_embeds, final_mask)即是第二节sam3_image.py中geo_feats, geo_masks的来历前者为几何 token 序列后者为其对应的 padding mask二者一起作为 prompt 参与后续 Transformer 的文本/几何联合编码。use_act_ckpt在构造时被保存但不在本模块内显式包装从源码结构看它由外部的封装层结合torch.utils.checkpoint机制统一启用。八、维度约定与使用要点速查综合Prompt、concat_padded_sequences与SequenceGeometryEncoder三者的 docstring 与断言可提炼出以下必须遵守的约定也是二次开发时最容易出错之处序列优先、批次第二embedding/坐标类张量形状为(seq_len, batch, feat)mask 与绝大多数标签张量批次优先。mask 的 1 表示 pad且必须是右填充所有交给编码器 / 拼接函数的 mask 都需满足is_right_padded。框坐标使用归一化 CxCyWH末维为 4转 xyxy、反归一化等由编码器内部按H/W完成外部只需保证归一化。d_model贯穿始终图像特征、位置编码输出、标签 embedding、投影输出与图像侧记忆共享同一维度改动时需保证 backbone 特征通道与d_model匹配。框与角点编码二选一encode_boxes_as_pointsTrue时内部忽略boxes_direct_project/boxes_pool/boxes_pos_enc三开关此时点编码仅用线性投影其标签字典扩大为 6 类。推理侧的最小化示例可以参考_get_dummy_promptpredict.py构造零框 Prompt 时传入box_embeddingstorch.zeros(0, B, 4)与box_masktorch.zeros(B, 0, dtypetorch.bool)即可编码器对空序列亦能正常前向直接拼接 CLS 后进入 Transformer。九、小结geometry_encoders.py是 SAM3 几何提示链路上封装最完整、契约最严格的模块之一Prompt统一了输入数据的形状与语义两个工具函数维护了变长序列批式运算的核心不变式SequenceGeometryEncoder则以多路编码求和 可选 CLS 后置精化 Transformer 堆叠的模块化设计把任意数量的归一化框提示转换为与图像/文本特征同构的 token 序列。本文所引的接口签名与行为均以 reference 文档、源码 及其在 build_sam3.py、sam3_image.py、predict.py 中的真实调用为准对为何如此设计等推断性结论已在文中以从源码结构看/可以推断等措辞明确标注。读者如需深入可直接以上述文件为入口研读完整实现。【免费下载链接】ultralyticsUltralytics YOLO26, YOLO11, YOLOv8 — object detection, instance segmentation, semantic segmentation, image classification, pose estimation, object tracking项目地址: https://gitcode.com/GitHub_Trending/ul/ultralytics创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考