资讯详情 解剖SAM推理链路:Ultralytics predict.py 源码实战指南
📅 2026/10/18 4:18:48
直接调ultralytics.models.sam.predict跑通一个分割 Demo 并不难难的是你想换一个输入模态、换一套后处理逻辑或者做性能优化时不知道去哪一层动刀。这篇文章我会把predict.py这个子模块拆开来讲先定位它在 SAM 推理链路里的真实身份再按__init__、preprocess、inference、postprocess的顺序走一遍最后把代码里不容易注意到的坑集中列出来。对刚接触源码的开发者或者想在业务里把 SAM 改造成工具的工程师这份阅读笔记应该能省你不少时间。1. 先定位predict.py 在 SAM 推理链路里的真实身份1.1 SAM 的“单次编码、多次提示”决定了代码结构SAM 不是传统意义上的一次性语义分割模型。它的核心设计是先用一个很重的图像编码器把整张图变成 image embedding再用一个相对轻量的 prompt encoder 和 mask decoder根据用户给出的点、框、掩码等提示信息输出分割结果。换句话说一张图只需要编码一次后续换不同提示去分割不同目标时重活不用重做。这个设计直接影响了predict.py的代码组织。SAMPredictor不需要像 YOLO 的预测器那样在inference里只做一次前向就完事它必须额外处理提示信息的组装、坐标归一化、多提示结果合并。所以你会发现这个文件里的核心方法并不是一个简单堆起来的模型调用而是一套围绕“图像 提示”的双输入数据流设计。阅读源码时心里要有这条主线图像走在一条通道里提示词走在另一条通道里两者在 mask decoder 汇合。predict.py的每个方法基本都在为这条汇合做准备。1.2 predict.py 干的是“填肉”不是“搭骨架”ultralytics里有一个很关键的基类叫BasePredictor它放在engine/predictor.py。这个基类已经把预测流程的骨架写好了读取source、循环取 batch、调用回调、保存结果、可视化输出。SAM 的predict.py继承了这个基类只重写其中几段关键逻辑。我一般把这种关系理解成“骨架 填肉”BasePredictor负责“循环”给定一批图片它知道怎么迭代。SAMPredictor负责“单个样本怎么处理”它决定输入如何预处理、模型如何被调用、输出如何转成统一结果对象。这也解释了为什么predict.py看起来方法不多。它不需要重新实现整个推理流程只需要精准覆盖 SAM 和 YOLO 的差异点。源码阅读最忌讳的是从头到尾逐行背诵正确做法是先搞清楚哪几处是子类自己实现的哪几处是继承来的。1.3 import 里的信息量打开ultralytics/models/sam/predict.py头部 import 一般能看到这几类from ultralytics.engine.predictor import BasePredictor from ultralytics.engine.results import Results from ultralytics.utils import DEFAULT_CFG, ops from ultralytics.utils.ops import non_max_suppression只看 import 就能推断出三件事这个子模块不是独立的它复用engine下的通用预测器。最后输出的是Results对象说明它和检测/分割任务共用一套结果规范。出现了non_max_suppression说明 SAM 的后处理里也有框层面的抑制逻辑不只是简单把 mask 拿出来。所以读predict.py时你的脑海里应该始终有一个更大的目录结构ultralytics/ └── models/ └── sam/ ├── build.py ├── predict.py ├── train.py └── modules/predict.py只是金字塔顶端的一层真正重的网络结构在modules/和build.py里。如果你想理解网络细节那里才是重点如果你想理解“这个模型怎么被产品化”predict.py反而是第一站。2. 从init开始SAMPredictor 的“默认值生意”2.1 构造方法里被强制钉死的三项配置SAMPredictor的__init__非常有辨识度开头通常会做一件事把默认参数强制覆盖成 SAM 需要的模式。简化后大致是这样的结构class SAMPredictor(BasePredictor): def __init__(self, cfgDEFAULT_CFG, overridesNone, _callbacksNone): if overrides is None: overrides {} overrides.update(conf0.25, tasksegment, modepredict) super().__init__(cfg, overrides, _callbacks) self.prompts {}这段代码的关键不在于super().__init__而在于overrides.update那三行。配置项默认值为什么钉死conf0.25过滤掉低置信度的 mask/box 预测是推理链路上的一个基础阈值tasksegment确保所有回调、保存逻辑和结果解析都走分割分支modepredict告诉整个框架这是推理模式不是训练或验证模式如果你是从 YOLO 检测任务转过来的你可能会觉得conf0.25是“检测框的置信度阈值”。但在 SAM 里它更多是 prompt decoder 输出的 mask 分数门槛。当你传了很多提示点模型会为每个提示组合生成候选 mask分数高低代表“这个 mask 和提示匹配的程度”。conf太低你会得到一堆碎片化的 maskconf太高有些提示点对应的目标可能直接被过滤掉。我最初读这里时有个误区以为重写__init__只是为了塞默认配置。后来才发现overrides.update是“给默认行为兜底”真正决定用户行为的是基类的参数合并规则。你要在业务里调低阈值直接在调用侧传conf0.1比硬改源码稳得多。2.2 setup_model 与 checkpoint 解析SAMPredictor 里另一个容易被忽略的方法是模型加载。YOLO 检测器通常用统一的AutoBackend读取各种权重格式但 SAM 不一样它需要build_sam根据 checkpoint 里记录的模型结构信息去组装模型。假如你传入的是sam_b.pt这类权重文件build_sam会解析出模型配置、加载权重、把模型搬上指定设备。这一步最大的坑是同一个.pt文件在不同版本下的兼容性不一定一致。我见过不少人拿老版本权重加载到新版本环境结果直接报结构不匹配。你在读predict.py时要把setup_model当作一个独立环节来看它不是为了简单读取权重而是为了把“权重文件”翻译成“可执行模型”。2.3 self.prompts一个容易被忽略的状态构造方法最后有一行self.prompts {}。这行代码非常不起眼但它承担了提示信息的状态管理。SAM 的推理模式里prompt 可以显式传入也可以暂存在预测器内部。self.prompts就是那个“暂存区”。实际调用时如果inference没有收到 prompts 参数预测器就会退回到这个内部变量取数据。这带来一个非常实际的问题如果你在一个共享的预测器实例上边跑边改self.prompts并发场景下会互相污染。如果是多线程处理图片我建议每个线程或协程维护自己的SAMPredictor实例而不是共用一个然后反复改内部状态。源码这么写是为了 API 方便但在生产环境里对共享状态要特别敏感。3. preprocess 与 inference数据从图像到 mask 的两级跳3.1 preprocess 到底预处理了什么preprocess方法接收的im通常已经是经过BasePredictor解析出来的图像数据。SAMPredictor里的preprocess只做“收尾工作”。我不建议把这段代码当成一个通用图像预处理模板来抄因为不同版本之间差异不小。稳定会出现的内容大致是这几件事把输入转成 Tensor搬到指定设备根据模型精度设置转成float或half检查张量维数不足 4 维就补一个 batch 维其中最容易忽视的是维数判断if len(im.shape) 4: im im.unsqueeze(0)你以为这只是“补一个 batch 维”其实它还有一层保护意义有些输入源解析出来是[C, H, W]有些是[H, W, C]如果上游没对齐这里先补维并不能解决通道顺序错误。所以在preprocess之前数据的通道顺序必须已经统一成模型需要的格式。很多人最后分割结果错位不是在postprocess找原因而是在preprocess这一步就埋下了隐患。3.2 inferenceprompt 是显式传进来的不是藏在模型里SAM 和传统模型一个很大的不同是它的inference方法通常接受两个输入图像和 prompts。代码骨架一般长这样def inference(self, im, promptsNone): if prompts is None: prompts self.prompts preds self.model(im, prompts) return preds你可能会觉得这是理所当然的但放在Ultralytics整个框架里看这是一个非常刻意的设计。因为BasePredictor的默认推理接口是不带 prompts 参数的子类如果要把 prompt 传下去就必须重新定义这个方法的签名。这就是子类“填肉”的典型体现。这个方法会暴露一个问题当你自己接入自定义模型时不一定能直接复用这段逻辑。比如你想把 SAM 的图像编码器换成自己的轻量网络或者你想在 prompt 进模型前做一次自定义编码最干净的方式是继承SAMPredictor重写inference而不是把模型内部的 forward 改得面目全非。3.3 点、框、mask 三种 prompt 的坐标约定predict.py本身不一定存放 prompt 编码的完整公式但它的接口层必须明确“prompt 长什么样”否则后续逻辑全部跑偏。类型常见传参形态说明pointpoints[[x, y], ...],labels[1, 0, ...]1 表示正例点0 表示负例/排除点boxboxes[[x1, y1, x2, y2], ...]目标框模型内部常转换为两个角点 promptmaskmasks数组通常是低分辨率或原图尺度的掩码用于迭代细化我把坐标约定单独拿出来说是因为它就是 SAM 落地的“隐形门槛”。外部 API 接收的往往是原始图像的像素坐标但模型内部为了统一处理通常会把坐标归一化到模型输入尺度。这个转换在哪做、按什么基准做决定了你在上游做图像缩放时会不会错位。举一个我真实踩过的场景为了加快预处理我先手动把图片resize到1024×1024再在这个缩放后的图片上选提示坐标。结果分割出来的目标总是偏的。原因很简单——提示坐标应该以原始图像为参考系而不是以我手动缩放后的网络输入为参考系。predict.py内部会把传进去的坐标按照它记录到的原始尺寸再做一次映射我提前缩放等于打乱了它的基准。所以在这类代码里最稳定的一条经验是不要在调用预测器之前自己偷偷改图像尺寸除非你把坐标同步改到同一个参考系。如果你确实需要小图提速也要把“坐标变换”和“图像变换”做成一对互逆操作。4. postprocess结果从模型输出到 Results 的最后一公里4.1 从 preds 拆出 mask、score、boxpostprocess是predict.py里最接近业务层的方法。它有三个入参preds、img、orig_imgs。preds是模型 forward 返回的原始结果img是网络输入张量orig_imgs是原始图片。这个方法第一件事往往是拆包。SAM 模型输出一般不会只有一个 mask通常会有 mask 预测、置信度或分数、以及和 prompt 对齐的 box 信息。在源码里你会看到类似这样的拆包逻辑pred_masks, pred_scores, pred_boxes preds[0], preds[1], preds[2]不同版本的下标可能不同我提醒自己每次都先打印一下preds的 shape再决定怎么拆。很多自定义改造翻车就是因为有人读文章看到这种拆包方式直接硬编码下标结果上游模型输出顺序一变后处理全错。4.2 NMS 和空结果保护predict.py里出现non_max_suppression说明即使 SAM 的输出主要是 mask它也要在“候选对象”层面做一次去重。因为当你传了多个 prompt模型可能为同一个目标生成多个高度重叠的 mask。做一次 NMS能保证最终返回的结果不会是一堆重复轮廓。真正让我觉得这段代码写得到位的地方是它处理“空结果”的方式。一般这种后处理代码里会有一个判断if pred_boxes.shape[0] 0: results.append(Results(orig_img, pathself.save_dir, namesnames)) continue意思是如果这一轮预测没有任何通过阈值过滤的候选框就不要制造一个格式错误而是返回一个空的Results对象。这样调用方的for result in results循环不会崩溃只是result.masks为空。这是工程上非常友好的设计。4.3 Results 对象才是和业务打交道的边界postprocess的最终目的是把张量数据封装成Results。这个对象不是简单字典而是一个带有orig_img、path、names、boxes、masks等字段的统一结果容器。它让下游代码可以这样写for result in results: print(result.orig_img.shape) print(result.masks.data.shape)这段代码的好处是无论上游是检测、分割还是 SAM调用侧的结构都一样。代价是调试时需要记一下字段的实际形状。比如result.masks.data通常是[N, H, W]而不是[N, 1, H, W]。我在刚开始接触时总会多打印一次 shape生怕在维度上翻车。5. 跑一个带提示点的实验调用方式反推代码逻辑5.1 最小复现脚本理论说完最好还是用一个最小脚本把源码里的逻辑串起来。以单点提示为例from ultralytics import SAM model SAM(sam_b.pt) results model.predict( sourcebus.jpg, points[[520, 220]], labels[1], saveTrue, projectsam_demo, ) for res in results: print(res.orig_img.shape) print(res.masks.data.shape)当你执行这段代码时内部流程大概是这样BasePredictor读取bus.jpg把它变成模型输入张量。SAMPredictor.preprocess负责补维和精度转换。SAMPredictor.inference拿到points和labels把图像张量和提示信息一起送进模型。模型返回一组 mask、score、box。SAMPredictor.postprocess过滤低分结果封装成Results。我在第一次跑通时最大的感受是predict.py把“提示信息怎么变成模型输入”这件事藏得太好了。你从外部 API 看不到self.prompts的完整结构只能靠读源码理解。所以如果你想做二次开发第一步不是改网络而是先打印中间变量比如在inference前后看一下 prompts 到底是什么形态。5.2 多 prompt 组合与负例点SAM 的 prompt 不限制只能给一个点。你可以在一次预测里传多个正例点和多个负例点results model.predict( sourcebus.jpg, points[[520, 220], [300, 180], [700, 350]], labels[1, 0, 1], )这里1表示这个点是目标区域0表示这个点不是目标区域。负例点特别适合用来排除和正例点靠得很近的干扰物体。比如你想分割前景里的人但旁边有棵树你可以用一个负例点落在树上让模型不要把它并进来。这段逻辑放到predict.py里看你会发现它的职责很清晰预测器本身不判断你的提示合不合理它只负责把提示对齐后送进模型。所以如果你传了一个超出图像边界的点或者points和labels长度不一致后处理阶段一定会出现各种难以理解的结果。5.3 三个常见报错和它们的本质原因现象本质原因建议提示点错位分割出来的目标偏了坐标参考系和模型输入不一致不要在上游手动 resize 图片或在 resize 时同步变换坐标报维度不对、广播失败points、labels、boxes的形状不符合接口约定统一用二维嵌套列表例如[[x, y]]不要写一维[x, y]去碰运气CPU 和半精度冲突模型被切到half但 CPU 算子不支持在 CPU 上强制使用float或在调用时明确设备为 GPU第二个错误尤其常见。很多人在传单个点时写points[520, 220]这是符合直觉的写法但和源码里期望的嵌套结构并不一致。源码解析时可能会按“多个点”去索引结果把一个坐标当成了一个点然后四舍五入得到完全不对的 mask。我自己的习惯是只要传点一律写[[x, y]]这种形式代码意图更清晰也不依赖接口做各种隐式兼容。6. 读完后我总结的坑缓存、精度与批量6.1 image embedding 没有透明的跨调用缓存SAM 最贵的部分在图像编码器。如果你拿着同一张图反复换提示点去试分割效果最理想的方案是图像只编码一次之后每次换 prompt 只跑轻量的 mask decoder。但从predict.py的代码结构来看它默认不会把这个 embedding 跨多次 predict 调用缓存为公开 API。也就是说如果这样写for point in point_list: model.predict(sourceimg_path, points[point], labels[1])那每次预测都会把整张图重新过一遍图像编码器。对交互式标注场景来说这个开销很浪费。想优化比较稳的做法是在业务层把“图像编码”和“提示预测”拆开先拿到图像 embedding 并缓存再在每次换提示时只做轻量解码。predict.py给你提供的是完整的推理闭环当你发现性能瓶颈在图像编码器时就得考虑绕过它而不是在子模块里塞一个状态量。6.2 半精度、CPU 与设备不一致SAM 这类模型在 GPU 上切half确实能省显存和加速。但如果你是在 CPU 上跑半精度经常会撞上一堆算子不支持的问题。你可能会看到类似“当前算子不支持 Half”的报错最后发现是模型精度状态和设备类型不匹配。predict.py的preprocess里会根据模型的fp16状态决定把输入转成half还是float。如果模型进了 CPU 但fp16还是开着问题往往在模型加载那里。调试这类问题第一步是确认model.device和model.fp16两个状态一致不要只在预处理里找原因。6.3 大 batch 和 prompt 批次的错位BasePredictor天然支持批量预测但 SAM 的 prompt 并不总是能像图像一样轻松对齐到 batch 维。如果你的调用方式是一张图配一组提示那么最简单、最不容易出错的做法是batch1。强行为了吞吐把四张图拼成一个 batch但每张图的提示个数又不一样代码复杂度会急剧上升。读predict.py时你会发现prompt 的“批”和图像的“批”并不是一回事。图像 batch 是模型输入维度的展开prompt 则更像一组并行的小样本。很多时候你以为加了大 batch 就能提速实测却发现不仅没提速反而因为提示补丁逻辑变复杂多花了不少功夫。这个坑在文档里不会写得实际跑一遍数据才看得清楚。6.4 改代码前先想清楚 override 还是复用最后说一个读源码的通用心得。看过predict.py之后你会觉得“这代码也不难我可以自己写一个”。但真正到业务改造时尽量优先继承SAMPredictor只重写你关心的那一段而不是把整个文件复制出去。比如你想自定义后处理就重写postprocess你想换提示编码方式就重写inference你只是想加一个回调直接复用BasePredictor的回调机制就行。predict.py的价值不在于每一行都写得精巧而在于它把“通用推理骨架”和“SAM 专属逻辑”切得非常干净。你只要找对那个切点就能用最小的改动盘活整个模型。我几次读回来之后最大的体会是像 SAM 这种提示驱动模型真正决定产品体验的不只是网络结构而是预测器怎么处理提示、坐标和结果封装。predict.py不花哨但它就是那个把模型从“学术代码”推向“工程落地”的关卡。