尧图网站建设 
简介面向深度学习与图像分割开发者的SAM2Segment Anything 2模型ONNX转换与推理Python工具包解决该先进分割模型跨框架部署难的问题适用于需要在ONNX Runtime、边缘设备或服务端快速落地实例分割任务的工程师和研究者。压缩包共14个文件以Python脚本、ONNX模型、示例图片和说明文档为主整体约591.77MB其中包含编码器与解码器onnx文件、推理脚本及测试图像可完整覆盖从模型导出到分割效果验证的流程。目前已有199人学习关注目录结构按主程序与模型文件划分便于对照理解。借助这份工具开发者无需自行搭建复杂的转换环境即可直接调用脚本将SAM2用于任意图像的实例分割也方便基于Python进行二次开发与参数调整显著降低部署门槛并提升图像处理项目的研发效率。1. ONNX-SAM2-Segment-Anything把 SAM2 从 PyTorch 搬到 ONNX Runtime 的落地脚本做图像分割落地的人应该都有同感SAM2 效果是好但 PyTorch 推理链路太重一台没有 GPU 的服务器根本跑不动更别说塞进边缘设备。这个 ONNX-SAM2-Segment-Anything 的 Python 脚本包解决的就是这个问题——它把 SAM2 的权重拆成图像编码器和 Mask 解码器两个 ONNX 模型用 onnxruntime 做推理CPU 上能跑量化成 int8 后也能往 RKNN、OpenVINO 这类推理栈上迁。对于要做图像分割集成、推理优化、模型迁移的开发者来说这套脚本的价值在于省掉了自己翻 SAM2 源码、改导出代码的大半时间。我拿到手之后先跑通了一遍流程再自己改了批量推理和量化校准下面把整个过程和踩过的坑完整拆给你。2. 为什么选 ONNX Runtime 跑 SAM2推理栈的取舍与脚本整体结构2.1 SAM2 模型拆解图像编码器和 Mask 解码器到底谁吃性能SAM2 的完整推理链路由三块组成图像编码器Image Encoder、Prompt 编码器Prompt Encoder和 Mask 解码器Mask Decoder。其中图像编码器采用的是 Hiera 结构的 ViT 变体输入一张 1024×1024 的 RGB 图输出的是多尺度的高分辨率特征图。这个编码器是整条链路里最重的一块无论是显存占用还是推理时延九成以上的成本都砸在它身上。Mask 解码器则很轻量它拿编码器输出的特征和用户给的点、框这类 prompt 做交叉注意力输出 1 到 4 个候选 mask。选择 ONNX Runtime 而不是继续用 PyTorch 推理核心原因有三个。第一onnxruntime 的 CPU 推理做了算子融合和线程池优化同样的 Hiera 编码器在无 GPU 环境下比 PyTorch 的 eager mode 快一到两倍这是一个可以直接实测的差距。第二ONNX 是中间表示导出之后可以无缝转到 TensorRT、OpenVINO、RKNN 这些平台一个模型文件多处部署。第三Python 侧只需要 onnxruntime 这一个推理依赖不需要把整个 torch 生态带进生产环境。常见做法是 CPU 环境装 onnxruntimeNVIDIA 显卡环境装 onnxruntime-gpu两者 API 一致代码几乎不用改。2.2 zip 包里的脚本分工先导出再推理两步走完整链路拿到 ONNX-SAM2-Segment-Anything.zip 之后先解压我一般会看一眼目录结构再动手。这套脚本的布局大致是这样文件/目录作用export_onnx.py从 PyTorch 权重导出 encoder 和 decoder 两个 ONNX 模型sam2_onnx_infer.py用 onnxruntime 加载 ONNX 做端到端分割推理preprocess.py图像缩放、归一化、padding 到 1024×1024requirements.txt依赖清单test_image/自带测试图用来验证导出和推理是否正常这种拆法的好处在于导出和推理解耦。你在带 GPU 的机器上导出一次 ONNX之后部署到任何环境都不再需要 PyTorch。推理脚本内部的主链路也很直接读图preprocess先跑 encoder 拿图像特征再把 prompt 坐标和特征一起喂给 decoder最后后处理输出 mask。导出和推理是两个独立阶段一旦遇到问题可以快速定位是模型导出有问题还是推理代码有问题。2.3 环境搭建onnxruntime 的 GPU 与 CPU 回退配置依赖安装是第一个容易翻车的地方。requirements.txt 里通常会有 onnxruntime 或 onnxruntime-gpu、onnx、opencv-python、numpy、pillow。我一般建议先确认目标机器的显卡和 CUDA 版本再决定装哪个。NVIDIA 显卡就装 onnxruntime-gpu版本要跟 CUDA 对应上纯 CPU 服务器装 onnxruntime 就行。这里有个坑onnxruntime-gpu 在没装 CUDA 的机器上也能 import 成功但一跑就报错报错信息还容易让人误以为是模型问题。安装命令是常规操作但值得写清楚# CPU 环境 pip install onnxruntime onnx opencv-python numpy pillow # GPU 环境注意版本与 CUDA 对应 pip install onnxruntime-gpu onnx opencv-python numpy pillow装完后可以用一小段代码验证推理后端是否正常加载import onnxruntime as ort # 打印可用执行提供程序GPU 环境会显示 CUDAExecutionProvider print(ort.get_available_providers()) session ort.InferenceSession( encoder.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider] )这段代码里providers参数的顺序是有讲究的。onnxruntime 会按列表顺序优先选择第一个可用的执行提供程序所以把 CUDA 放前面、CPU 放后面就实现了「有 GPU 用 GPU没 GPU 自动回退 CPU」的部署策略。很多人在部署阶段喜欢写死 provider结果换台机器就崩这个习惯不好。我一般都会列出两个 provider 做兜底部署省心得多。3. PyTorch 转 ONNX图像编码器与 Mask 解码器的导出细节3.1 导出前必须搞懂的几个关键配置PyTorch 转 ONNX 不是一句torch.onnx.export就完事的。SAM2 这类视觉模型有几个特殊点不处理清楚导出来的 ONNX 要么跑不通要么输出错得离谱。第一个是opset版本。SAM2 的源码里用到了很多较新的算子比如torch.nn.functional.scaled_dot_product_attention这类操作在低版本 opset 里没有对应的映射。我实测下来opset 至少要到 17建议直接设 17 或 18。设低了导出时就会报「不支持的算子」设高了又要求目标环境的 onnxruntime 版本足够新。这个平衡要在导出时就定好。第二个是动态轴dynamic_axes。图像编码器的输入是 1024×1024如果你写死静态 shape导出的 ONNX 就只能跑这一个分辨率换一张图就报维度错误。所以导出时要把 batch、高度、宽度都设成动态轴。第三个是do_constant_folding。这个参数建议打开它会在导出阶段把能算的常量提前算掉推理时少跑一部分节点。对 Hiera 这种带大量位置编码和相对位置索引的模型收益挺明显。3.2 图像编码器导出动态分辨率与 batch 设置下面是脚本里导出图像编码器的核心代码。SAM2 的权重是官方 PyTorch checkpoint先加载再拆出 encoder 部分做导出import torch from sam2.modeling.sam2_base import SAM2Base from sam2.build_sam import build_sam2 # 加载官方权重配置文件按实际权重选择 sam2 build_sam2( config_filesam2_hiera_b.yaml, ckpt_pathsam2_hiera_base.pt, devicecuda ) # 取出图像编码器部分转为 eval 模式 image_encoder sam2.image_encoder.eval() # 构造示例输入1 张 1024x1024 的图 dummy_input torch.randn(1, 3, 1024, 1024, devicecuda) # 导出 ONNX关键参数见下方说明 torch.onnx.export( image_encoder, dummy_input, encoder.onnx, input_names[images], output_names[image_features], dynamic_axes{ images: {0: batch, 2: height, 3: width}, image_features: {0: batch, 2: height, 3: width} }, opset_version17, do_constant_foldingTrue )这段代码里有几个参数值得单独拎出来说。input_names和output_names不是随便起的它们会成为 ONNX 图里的输入输出节点名后续推理脚本要用同样的名字喂数据。dynamic_axes里的{0: batch, 2: height, 3: width}表示这三个维度在推理时可以变化维度名字叫 batch、height、width。这样导出的模型既可以跑单张图也可以一次喂四张图批量推理。opset 设 17 是为了兼容注意力算子的导出。如果导出时报算子不支持优先检查 torch 版本和 opset 版本是否匹配。导出完成后建议立即检查一下 ONNX 模型有效性import onnx model onnx.load(encoder.onnx) onnx.checker.check_model(model) print(encoder.onnx valid)这一步不花多少时间但能提前拦掉一批「导出时没报错、推理时崩」的问题。3.3 Mask 解码器导出多输入多输出的动态轴处理Mask 解码器比编码器麻烦因为它的输入输出都很多。输入包括图像特征、提示点的坐标、提示点的标签、还有可选的框输入输出包括多个候选 mask 和对应的置信度分数。导出的时候每个输入都要单独命名动态轴也要逐个处理# 构造 decoder 的示例输入 dummy_feats torch.randn(1, 256, 64, 64, devicecuda) dummy_point_coords torch.randn(1, 5, 2, devicecuda) dummy_point_labels torch.zeros(1, 5, dtypetorch.int64, devicecuda) dummy_box torch.randn(1, 4, devicecuda) torch.onnx.export( sam2.mask_decoder, (dummy_feats, dummy_point_coords, dummy_point_labels, dummy_box), decoder.onnx, input_names[image_features, point_coords, point_labels, box], output_names[masks, iou_predictions], dynamic_axes{ image_features: {0: batch, 2: height, 3: width}, point_coords: {0: batch, 1: num_points}, point_labels: {0: batch, 1: num_points}, masks: {0: batch, 2: height, 3: width} }, opset_version17, do_constant_foldingTrue )这里输入的维度含义要清楚point_coords的形状是(batch, num_points, 2)最后一位分别存 x 和 y 坐标point_labels是(batch, num_points)1 表示前景点0 表示背景点这个标签体系不能改改了分割结果就是反的。box输入是(batch, 4)顺序是 x1, y1, x2, y2是绝对坐标而不是归一化坐标。我见过不少人在这里栽跟头——prompt 坐标本身是对的但标签给反了出来的 mask 整个反转。3.4 导出后验证onnxruntime 输出与 PyTorch 输出的误差对齐导出不等于成功还得验证 ONNX 的输出跟 PyTorch 原模型差多少。这一步是必要的因为有些算子导出后在数值上会有微小漂移如果不做对齐就直接量化部署后面出了问题根本分不清是转换的锅还是量化的锅。import onnxruntime as ort import numpy as np # PyTorch 参考输出 with torch.no_grad(): ref_feats image_encoder(dummy_input).cpu().numpy() # ONNX Runtime 推理 ort_session ort.InferenceSession( encoder.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider] ) ort_out ort_session.run( [image_features], {images: dummy_input.cpu().numpy()} )[0] # 计算最大绝对误差 max_diff np.max(np.abs(ref_feats - ort_out)) print(fmax abs diff: {max_diff:.6f})误差在 1e-3 量级以内就算正常说明导出链路没问题如果误差到 1e-1 甚至出现 NaN基本可以断定是某个算子映射出了问题要回头检查 torch 和 onnxruntime 的版本。我在实际项目中用这套验证方法成功定位过两次问题一次是 torch 版本太老导致scaled_dot_product_attention导出成了低精度版本一次是输入张量没有转成 float32 导致精度截断。这类问题在导出阶段拦下来成本最低。4. INT8 量化与精度诊断脚本自带的量化参数和调优手段4.1 哪个模型值得量化Encoder 量化收益大但风险更高部署场景里模型体积和推理速度经常比精度更重要。SAM2 的 encoder 是 Hiera 结构FP32 模型动辄几百 MB在边缘设备上根本放不下decoder 却很小一般只有几 MB。所以量化策略很明确重点压 encoderdecoder 保持 FP32 甚至直接不量化。听起来简单但 encoder 量化恰恰是风险最高的部分。Hiera 的注意力模块对量化误差敏感尤其是 softmax 之后的特征分布一旦被 int8 截断会导致编码器输出的特征图整体漂移。这个漂移不像分类模型那样只影响置信度它会直接让 decode 出来的 mask 形状破裂。我一般会做一个中间态方案只量化 encoder 里的卷积层和全连接层把注意力部分的算子保留为 FP32。ONNX Runtime 的静态量化支持按算子粒度设置排除表可以做到这一点。这个方案的体积压缩率低一些但精度损失能控制在很小的范围内。完整量化是锦上添花先跑通混合精度才是稳妥路径。4.2 静态量化与校准数据集校准集不能随便选onnxruntime 的静态量化需要一段校准过程用真实数据统计每个激活张量的数值范围。脚本量化部分用到的就是quantize_static。完整流程是先跑一遍校准数据收集再做量化最后对比精度。核心代码如下from onnxruntime.quantization import quantize_static, QuantType from onnxruntime.quantization.calibrate import CalibrationDataReader class SAM2CalibDataReader(CalibrationDataReader): def __init__(self, images): self.images images self.idx 0 def get_next(self): if self.idx len(self.images): return None data {images: self.images[self.idx]} self.idx 1 return data def rewind(self): self.idx 0 # 校准数据必须来自真实业务场景数量不能太少 calib_images load_100_business_images() quantize_static( model_inputencoder.onnx, model_outputencoder_int8.onnx, calibration_data_readerSAM2CalibDataReader(calib_images), quant_formatQuantFormat.QDQ, per_channelTrue, weight_typeQuantType.QInt8, activation_typeQuantType.QUInt8, extra_options{ActivationSymmetric: False} )这段代码里的几个选项直接决定量化质量。per_channelTrue表示每个输出通道单独计算缩放系数精度比 per-tensor 高不少强烈建议打开。ActivationSymmetricFalse表示激活值使用非对称量化对 ReLU 这类非负分布效果好。quant_format选QDQ格式它的好处是可以部分算子回退 FP32混合精度方案就是基于这种格式。校准数据这块我踩过的坑是图省事用了 COCO 数据集做校准结果在工业质检场景里 quant 出来的模型 mask 全是破洞——校准集的分布离真实场景太远统计出来的数值范围是错的。后来换成 100 张真实产线图问题才解决。校准数据的数量经验值是 50 到 200 张。少于 50 张统计出来的 min/max 不稳定多于 200 张耗时增长明显但精度提升趋近于零。4.3 量化后精度对比mask 的 IoU 掉多少才算可接受量化完不能直接用必须做量化前后精度对比。对分割模型来说最直观的指标是 mask IoU。我常用的验证方法是固定一组 prompt分别用 FP32 和 int8 模型跑同一个 batch计算两者 mask 输出的 IoU 或 Dicefrom sklearn.metrics import jaccard_score def compute_iou(fp32_mask, int8_mask): # 两个 mask 都已经是 0/1 的二值图 intersect np.logical_and(fp32_mask, int8_mask).sum() union np.logical_or(fp32_mask, int8_mask).sum() return intersect / (union 1e-6) iou_values [] for fp32_mask, int8_mask in zip(fp32_masks, int8_masks): iou_values.append(compute_iou(fp32_mask, int8_mask)) mean_iou np.mean(iou_values) print(fFP32 vs INT8 mIoU: {mean_iou:.4f})这个指标的含义是int8 输出与 fp32 输出在像素级上的重合度。实操中mIoU 在 0.97 以上算优秀0.95 到 0.97 属于正常波动低于 0.93 就说明量化方案需要调整。如果量化把编码器特征的空间细节磨掉了decoder 出来的 mask 边缘会明显粗糙最直接的修正方式是给 encoder 的注意力算子做 FP32 回退。量化精度这块一直被戏称为玄学因为同样的参数换一批校准数据结果能差好几个点。但本质上还是数值分布统计的问题只要校准集贴近真实场景、per-channel 打开、必要时做混合精度结果通常都能落在可用范围内。5. 常见问题排查这个脚本跑不通时的五个翻车点5.1 解压就报错CRC 校验失败或提示文件损坏现象解压 ONNX-SAM2-Segment-Anything.zip 时提示 CRC 失败甚至直接说文件损坏脚本根本拿不出来。原因下载过程中文件不完整或者 zip 包本身带了一些在老旧解压工具下不兼容的压缩特性。还有一种情况是文件在传输时被某些安全软件拦截改写。解决先核对文件大小是否和下载页面标注一致不一致就重新下载。如果文件大小没问题换 7-Zip 或者 WinRAR 的高版本重试这类工具对 zip 扩展特性的兼容性比系统自带的好。我自己处理过的一个案例是用系统自带的解压一直提示伪加密换 7-Zip 之后正常解出。解压之后最好直接核对一下文件数防止中途缺文件。5.2 Session 创建失败提示算子不存在或 provider 不可用现象ort.InferenceSession创建时抛异常报错信息里出现「Unsupported Operator」或「Execution Provider not found」。原因ONNX 模型的 opset 版本太高当前 onnxruntime 版本不支持或者装的是 CPU 版 onnxruntime但代码里强行指定了 CUDAExecutionProvider。解决优先检查 onnxruntime 版本。如果模型是 opset 18 导出的onnxruntime 至少要 1.16 以上。其次检查 provider 指定逻辑不要写死单一 provider用列表同时声明 CUDA 和 CPU 做回退。这类问题在换机器部署时最容易冒出来源头大都是导出时 opset 设得过于激进我现在的习惯是 opset 统一设 17兼容性最好。5.3 推理结果全是黑图或白图mask 输出异常现象代码正常跑完但输出的 mask 是全黑或者全白看起来像是分割失败但没有任何报错。原因两种典型情况。第一种是输入图像没有做归一化FP32 的原始像素范围是 0 到 255而模型训练时用的是 0 到 1 的归一化值特征分布完全错位。第二种是 prompt 坐标体系不对SAM2 要求 prompt 坐标是输入图 resize 后的绝对坐标如果直接传原图坐标prompt 点落在完全不同的位置上分割结果自然会乱掉。解决检查预处理环节确认图像在送入 encoder 前已经缩放到 1024×1024 并做了归一化prompt 坐标要做同比例的换算。这个环节是黑匣子出错时不报错、不给提示只能靠输出结果反推所以我把 preprocess 和坐标换算单独抽出来做了单元测试。5.4 量化后 mask 边缘像锯齿精度掉了但不报错现象int8 模型跑出来的 mask 形状基本对但边缘全是锯齿小目标直接消失。原因encoder 量化粗粒度导致的特征精度损失。calibration 数据集太少或者分布偏统计的激活值范围不准导致量化误差集中在细节特征上注意力算子的量化损失也会放大边缘噪声。解决回到量化环节增加校准数据到 100 张以上检查 per_channel 是否打开如果还不行就做混合精度把注意力算子的量化排除掉。量化效果看 mIoU 而不是肉眼看图这是血泪经验——有时候肉眼看着差不多但实际上已经掉了三个点。5.5 内存溢出批量推理时图像一次性全载入现象批量推理时提示 OOM 或 MemoryError尤其是带 GPU 的环境更容易触发。原因onnxruntime 的 session 会缓存部分中间张量在多 batch 和大分辨率下显存占用会叠加。特别是 encoder 的输出特征图是四维浮点张量批量一大就超出了显存限额。解决控制推理 batch 大小不要一次性塞几十张图。常见做法是每次喂 4 张或 8 张循环处理。还有一个技巧是复用 session 的run时尽量用同一个输入缓冲区减少张量重新分配的频率。6. 批量分割与结果验证把脚本改造成可落地的推理管线脚本本身跑通之后离生产落地还差一步批量处理多张图、保存结果、量化质量验证。这一节我给出一个可以直接用的批量推理改造版也是我把这套脚本接到实际项目里时做的改动。批量推理的核心是对每张图调用一轮「预处理 → 编码器 → 解码器 → 后处理」把 mask 存成 PNG 文件。prompt 部分我固定用中心点作为前景提示适合「对每张图的主体做分割」这类场景如果你要指定框或多个点改成从表格或 JSON 读取即可import cv2 import glob import numpy as np import onnxruntime as ort # 推理 session 全局复用避免反复创建 enc_session ort.InferenceSession(encoder_int8.onnx, providers[CPUExecutionProvider]) dec_session ort.InferenceSession(decoder.onnx, providers[CPUExecutionProvider]) def infer_one(image_path, out_path, pointNone): img cv2.imread(image_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 预处理缩放、归一化、转 CHW img_resized cv2.resize(img, (1024, 1024)) x img_resized.astype(np.float32) / 255.0 x np.transpose(x, (2, 0, 1))[None] # 变成 1x3x1024x1024 feats enc_session.run([image_features], {images: x})[0] # prompt 默认中心点标签为前景 h, w 1024, 1024 if point is None: point_coords np.array([[[w // 2, h // 2]]], dtypenp.float32) else: point_coords np.array([[point]], dtypenp.float32) point_labels np.ones((1, 1), dtypenp.int64) box np.zeros((1, 4), dtypenp.float32) # 空框不用框提示 masks, ious dec_session.run( [masks, iou_predictions], { image_features: feats, point_coords: point_coords, point_labels: point_labels, box: box } ) # SAM2 输出多个候选 mask取 IoU 最高的那个 best np.argmax(ious[0]) mask masks[0][best] mask (mask 0.0).astype(np.uint8) * 255 cv2.imwrite(out_path, mask) for img_path in glob.glob(images/*.jpg): out_path img_path.replace(images, outputs) infer_one(img_path, out_path, pointNone)这段代码里有三个细节值得说明。第一encoder 和 decoder 两个 session 各自复用不要在循环里反复创建否则光是初始化开销就能让整体耗时翻倍。第二masks的输出形状是(batch, num_masks, H, W)iou_predictions是(batch, num_masks)取argmax选最高置信度的 mask 是 SAM 系列的标准后处理方式。第三mask 阈值用了 0.0因为 SAM2 的 mask 输出已经带 sigmoid 激活大于 0.5 会更严格0.0 能保留更多边缘细节高强度场景可以自己调。批量处理跑通之后我还做了一次全量验证对 200 张测试图分别用 FP32 模型和 int8 模型跑分割计算两版 mask 的 mIoU。那次量化我一开始只用了 10 张校准图结果 mIoU 掉到 0.90边缘大面积破损后来把校准图加到 120 张、打开 per-channelmIoU 回到了 0.96边缘细节基本恢复。从那以后我每次量化都强制走一遍「先校准、再验证、最后看 mIoU」的流程不再凭感觉判断量化好坏。这套流程看起来多花半小时但能省掉后面部署现场排查问题的一整天。希望帮到你。本文还有配套的精品资源点击获取 SAM2 PyTorch 转 ONNX Runtime 部署与 INT8 量化
简介面向深度学习与图像分割开发者的SAM2Segment Anything 2模型ONNX转换与推理Python工具包解决该先进分割模型跨框架部署难的问题适用于需要在ONNX Runtime、边缘设备或服务端快速落地实例分割任务的工程师和研究者。压缩包共14个文件以Python脚本、ONNX模型、示例图片和说明文档为主整体约591.77MB其中包含编码器与解码器onnx文件、推理脚本及测试图像可完整覆盖从模型导出到分割效果验证的流程。目前已有199人学习关注目录结构按主程序与模型文件划分便于对照理解。借助这份工具开发者无需自行搭建复杂的转换环境即可直接调用脚本将SAM2用于任意图像的实例分割也方便基于Python进行二次开发与参数调整显著降低部署门槛并提升图像处理项目的研发效率。1. ONNX-SAM2-Segment-Anything把 SAM2 从 PyTorch 搬到 ONNX Runtime 的落地脚本做图像分割落地的人应该都有同感SAM2 效果是好但 PyTorch 推理链路太重一台没有 GPU 的服务器根本跑不动更别说塞进边缘设备。这个 ONNX-SAM2-Segment-Anything 的 Python 脚本包解决的就是这个问题——它把 SAM2 的权重拆成图像编码器和 Mask 解码器两个 ONNX 模型用 onnxruntime 做推理CPU 上能跑量化成 int8 后也能往 RKNN、OpenVINO 这类推理栈上迁。对于要做图像分割集成、推理优化、模型迁移的开发者来说这套脚本的价值在于省掉了自己翻 SAM2 源码、改导出代码的大半时间。我拿到手之后先跑通了一遍流程再自己改了批量推理和量化校准下面把整个过程和踩过的坑完整拆给你。2. 为什么选 ONNX Runtime 跑 SAM2推理栈的取舍与脚本整体结构2.1 SAM2 模型拆解图像编码器和 Mask 解码器到底谁吃性能SAM2 的完整推理链路由三块组成图像编码器Image Encoder、Prompt 编码器Prompt Encoder和 Mask 解码器Mask Decoder。其中图像编码器采用的是 Hiera 结构的 ViT 变体输入一张 1024×1024 的 RGB 图输出的是多尺度的高分辨率特征图。这个编码器是整条链路里最重的一块无论是显存占用还是推理时延九成以上的成本都砸在它身上。Mask 解码器则很轻量它拿编码器输出的特征和用户给的点、框这类 prompt 做交叉注意力输出 1 到 4 个候选 mask。选择 ONNX Runtime 而不是继续用 PyTorch 推理核心原因有三个。第一onnxruntime 的 CPU 推理做了算子融合和线程池优化同样的 Hiera 编码器在无 GPU 环境下比 PyTorch 的 eager mode 快一到两倍这是一个可以直接实测的差距。第二ONNX 是中间表示导出之后可以无缝转到 TensorRT、OpenVINO、RKNN 这些平台一个模型文件多处部署。第三Python 侧只需要 onnxruntime 这一个推理依赖不需要把整个 torch 生态带进生产环境。常见做法是 CPU 环境装 onnxruntimeNVIDIA 显卡环境装 onnxruntime-gpu两者 API 一致代码几乎不用改。2.2 zip 包里的脚本分工先导出再推理两步走完整链路拿到 ONNX-SAM2-Segment-Anything.zip 之后先解压我一般会看一眼目录结构再动手。这套脚本的布局大致是这样文件/目录作用export_onnx.py从 PyTorch 权重导出 encoder 和 decoder 两个 ONNX 模型sam2_onnx_infer.py用 onnxruntime 加载 ONNX 做端到端分割推理preprocess.py图像缩放、归一化、padding 到 1024×1024requirements.txt依赖清单test_image/自带测试图用来验证导出和推理是否正常这种拆法的好处在于导出和推理解耦。你在带 GPU 的机器上导出一次 ONNX之后部署到任何环境都不再需要 PyTorch。推理脚本内部的主链路也很直接读图preprocess先跑 encoder 拿图像特征再把 prompt 坐标和特征一起喂给 decoder最后后处理输出 mask。导出和推理是两个独立阶段一旦遇到问题可以快速定位是模型导出有问题还是推理代码有问题。2.3 环境搭建onnxruntime 的 GPU 与 CPU 回退配置依赖安装是第一个容易翻车的地方。requirements.txt 里通常会有 onnxruntime 或 onnxruntime-gpu、onnx、opencv-python、numpy、pillow。我一般建议先确认目标机器的显卡和 CUDA 版本再决定装哪个。NVIDIA 显卡就装 onnxruntime-gpu版本要跟 CUDA 对应上纯 CPU 服务器装 onnxruntime 就行。这里有个坑onnxruntime-gpu 在没装 CUDA 的机器上也能 import 成功但一跑就报错报错信息还容易让人误以为是模型问题。安装命令是常规操作但值得写清楚# CPU 环境 pip install onnxruntime onnx opencv-python numpy pillow # GPU 环境注意版本与 CUDA 对应 pip install onnxruntime-gpu onnx opencv-python numpy pillow装完后可以用一小段代码验证推理后端是否正常加载import onnxruntime as ort # 打印可用执行提供程序GPU 环境会显示 CUDAExecutionProvider print(ort.get_available_providers()) session ort.InferenceSession( encoder.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider] )这段代码里providers参数的顺序是有讲究的。onnxruntime 会按列表顺序优先选择第一个可用的执行提供程序所以把 CUDA 放前面、CPU 放后面就实现了「有 GPU 用 GPU没 GPU 自动回退 CPU」的部署策略。很多人在部署阶段喜欢写死 provider结果换台机器就崩这个习惯不好。我一般都会列出两个 provider 做兜底部署省心得多。3. PyTorch 转 ONNX图像编码器与 Mask 解码器的导出细节3.1 导出前必须搞懂的几个关键配置PyTorch 转 ONNX 不是一句torch.onnx.export就完事的。SAM2 这类视觉模型有几个特殊点不处理清楚导出来的 ONNX 要么跑不通要么输出错得离谱。第一个是opset版本。SAM2 的源码里用到了很多较新的算子比如torch.nn.functional.scaled_dot_product_attention这类操作在低版本 opset 里没有对应的映射。我实测下来opset 至少要到 17建议直接设 17 或 18。设低了导出时就会报「不支持的算子」设高了又要求目标环境的 onnxruntime 版本足够新。这个平衡要在导出时就定好。第二个是动态轴dynamic_axes。图像编码器的输入是 1024×1024如果你写死静态 shape导出的 ONNX 就只能跑这一个分辨率换一张图就报维度错误。所以导出时要把 batch、高度、宽度都设成动态轴。第三个是do_constant_folding。这个参数建议打开它会在导出阶段把能算的常量提前算掉推理时少跑一部分节点。对 Hiera 这种带大量位置编码和相对位置索引的模型收益挺明显。3.2 图像编码器导出动态分辨率与 batch 设置下面是脚本里导出图像编码器的核心代码。SAM2 的权重是官方 PyTorch checkpoint先加载再拆出 encoder 部分做导出import torch from sam2.modeling.sam2_base import SAM2Base from sam2.build_sam import build_sam2 # 加载官方权重配置文件按实际权重选择 sam2 build_sam2( config_filesam2_hiera_b.yaml, ckpt_pathsam2_hiera_base.pt, devicecuda ) # 取出图像编码器部分转为 eval 模式 image_encoder sam2.image_encoder.eval() # 构造示例输入1 张 1024x1024 的图 dummy_input torch.randn(1, 3, 1024, 1024, devicecuda) # 导出 ONNX关键参数见下方说明 torch.onnx.export( image_encoder, dummy_input, encoder.onnx, input_names[images], output_names[image_features], dynamic_axes{ images: {0: batch, 2: height, 3: width}, image_features: {0: batch, 2: height, 3: width} }, opset_version17, do_constant_foldingTrue )这段代码里有几个参数值得单独拎出来说。input_names和output_names不是随便起的它们会成为 ONNX 图里的输入输出节点名后续推理脚本要用同样的名字喂数据。dynamic_axes里的{0: batch, 2: height, 3: width}表示这三个维度在推理时可以变化维度名字叫 batch、height、width。这样导出的模型既可以跑单张图也可以一次喂四张图批量推理。opset 设 17 是为了兼容注意力算子的导出。如果导出时报算子不支持优先检查 torch 版本和 opset 版本是否匹配。导出完成后建议立即检查一下 ONNX 模型有效性import onnx model onnx.load(encoder.onnx) onnx.checker.check_model(model) print(encoder.onnx valid)这一步不花多少时间但能提前拦掉一批「导出时没报错、推理时崩」的问题。3.3 Mask 解码器导出多输入多输出的动态轴处理Mask 解码器比编码器麻烦因为它的输入输出都很多。输入包括图像特征、提示点的坐标、提示点的标签、还有可选的框输入输出包括多个候选 mask 和对应的置信度分数。导出的时候每个输入都要单独命名动态轴也要逐个处理# 构造 decoder 的示例输入 dummy_feats torch.randn(1, 256, 64, 64, devicecuda) dummy_point_coords torch.randn(1, 5, 2, devicecuda) dummy_point_labels torch.zeros(1, 5, dtypetorch.int64, devicecuda) dummy_box torch.randn(1, 4, devicecuda) torch.onnx.export( sam2.mask_decoder, (dummy_feats, dummy_point_coords, dummy_point_labels, dummy_box), decoder.onnx, input_names[image_features, point_coords, point_labels, box], output_names[masks, iou_predictions], dynamic_axes{ image_features: {0: batch, 2: height, 3: width}, point_coords: {0: batch, 1: num_points}, point_labels: {0: batch, 1: num_points}, masks: {0: batch, 2: height, 3: width} }, opset_version17, do_constant_foldingTrue )这里输入的维度含义要清楚point_coords的形状是(batch, num_points, 2)最后一位分别存 x 和 y 坐标point_labels是(batch, num_points)1 表示前景点0 表示背景点这个标签体系不能改改了分割结果就是反的。box输入是(batch, 4)顺序是 x1, y1, x2, y2是绝对坐标而不是归一化坐标。我见过不少人在这里栽跟头——prompt 坐标本身是对的但标签给反了出来的 mask 整个反转。3.4 导出后验证onnxruntime 输出与 PyTorch 输出的误差对齐导出不等于成功还得验证 ONNX 的输出跟 PyTorch 原模型差多少。这一步是必要的因为有些算子导出后在数值上会有微小漂移如果不做对齐就直接量化部署后面出了问题根本分不清是转换的锅还是量化的锅。import onnxruntime as ort import numpy as np # PyTorch 参考输出 with torch.no_grad(): ref_feats image_encoder(dummy_input).cpu().numpy() # ONNX Runtime 推理 ort_session ort.InferenceSession( encoder.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider] ) ort_out ort_session.run( [image_features], {images: dummy_input.cpu().numpy()} )[0] # 计算最大绝对误差 max_diff np.max(np.abs(ref_feats - ort_out)) print(fmax abs diff: {max_diff:.6f})误差在 1e-3 量级以内就算正常说明导出链路没问题如果误差到 1e-1 甚至出现 NaN基本可以断定是某个算子映射出了问题要回头检查 torch 和 onnxruntime 的版本。我在实际项目中用这套验证方法成功定位过两次问题一次是 torch 版本太老导致scaled_dot_product_attention导出成了低精度版本一次是输入张量没有转成 float32 导致精度截断。这类问题在导出阶段拦下来成本最低。4. INT8 量化与精度诊断脚本自带的量化参数和调优手段4.1 哪个模型值得量化Encoder 量化收益大但风险更高部署场景里模型体积和推理速度经常比精度更重要。SAM2 的 encoder 是 Hiera 结构FP32 模型动辄几百 MB在边缘设备上根本放不下decoder 却很小一般只有几 MB。所以量化策略很明确重点压 encoderdecoder 保持 FP32 甚至直接不量化。听起来简单但 encoder 量化恰恰是风险最高的部分。Hiera 的注意力模块对量化误差敏感尤其是 softmax 之后的特征分布一旦被 int8 截断会导致编码器输出的特征图整体漂移。这个漂移不像分类模型那样只影响置信度它会直接让 decode 出来的 mask 形状破裂。我一般会做一个中间态方案只量化 encoder 里的卷积层和全连接层把注意力部分的算子保留为 FP32。ONNX Runtime 的静态量化支持按算子粒度设置排除表可以做到这一点。这个方案的体积压缩率低一些但精度损失能控制在很小的范围内。完整量化是锦上添花先跑通混合精度才是稳妥路径。4.2 静态量化与校准数据集校准集不能随便选onnxruntime 的静态量化需要一段校准过程用真实数据统计每个激活张量的数值范围。脚本量化部分用到的就是quantize_static。完整流程是先跑一遍校准数据收集再做量化最后对比精度。核心代码如下from onnxruntime.quantization import quantize_static, QuantType from onnxruntime.quantization.calibrate import CalibrationDataReader class SAM2CalibDataReader(CalibrationDataReader): def __init__(self, images): self.images images self.idx 0 def get_next(self): if self.idx len(self.images): return None data {images: self.images[self.idx]} self.idx 1 return data def rewind(self): self.idx 0 # 校准数据必须来自真实业务场景数量不能太少 calib_images load_100_business_images() quantize_static( model_inputencoder.onnx, model_outputencoder_int8.onnx, calibration_data_readerSAM2CalibDataReader(calib_images), quant_formatQuantFormat.QDQ, per_channelTrue, weight_typeQuantType.QInt8, activation_typeQuantType.QUInt8, extra_options{ActivationSymmetric: False} )这段代码里的几个选项直接决定量化质量。per_channelTrue表示每个输出通道单独计算缩放系数精度比 per-tensor 高不少强烈建议打开。ActivationSymmetricFalse表示激活值使用非对称量化对 ReLU 这类非负分布效果好。quant_format选QDQ格式它的好处是可以部分算子回退 FP32混合精度方案就是基于这种格式。校准数据这块我踩过的坑是图省事用了 COCO 数据集做校准结果在工业质检场景里 quant 出来的模型 mask 全是破洞——校准集的分布离真实场景太远统计出来的数值范围是错的。后来换成 100 张真实产线图问题才解决。校准数据的数量经验值是 50 到 200 张。少于 50 张统计出来的 min/max 不稳定多于 200 张耗时增长明显但精度提升趋近于零。4.3 量化后精度对比mask 的 IoU 掉多少才算可接受量化完不能直接用必须做量化前后精度对比。对分割模型来说最直观的指标是 mask IoU。我常用的验证方法是固定一组 prompt分别用 FP32 和 int8 模型跑同一个 batch计算两者 mask 输出的 IoU 或 Dicefrom sklearn.metrics import jaccard_score def compute_iou(fp32_mask, int8_mask): # 两个 mask 都已经是 0/1 的二值图 intersect np.logical_and(fp32_mask, int8_mask).sum() union np.logical_or(fp32_mask, int8_mask).sum() return intersect / (union 1e-6) iou_values [] for fp32_mask, int8_mask in zip(fp32_masks, int8_masks): iou_values.append(compute_iou(fp32_mask, int8_mask)) mean_iou np.mean(iou_values) print(fFP32 vs INT8 mIoU: {mean_iou:.4f})这个指标的含义是int8 输出与 fp32 输出在像素级上的重合度。实操中mIoU 在 0.97 以上算优秀0.95 到 0.97 属于正常波动低于 0.93 就说明量化方案需要调整。如果量化把编码器特征的空间细节磨掉了decoder 出来的 mask 边缘会明显粗糙最直接的修正方式是给 encoder 的注意力算子做 FP32 回退。量化精度这块一直被戏称为玄学因为同样的参数换一批校准数据结果能差好几个点。但本质上还是数值分布统计的问题只要校准集贴近真实场景、per-channel 打开、必要时做混合精度结果通常都能落在可用范围内。5. 常见问题排查这个脚本跑不通时的五个翻车点5.1 解压就报错CRC 校验失败或提示文件损坏现象解压 ONNX-SAM2-Segment-Anything.zip 时提示 CRC 失败甚至直接说文件损坏脚本根本拿不出来。原因下载过程中文件不完整或者 zip 包本身带了一些在老旧解压工具下不兼容的压缩特性。还有一种情况是文件在传输时被某些安全软件拦截改写。解决先核对文件大小是否和下载页面标注一致不一致就重新下载。如果文件大小没问题换 7-Zip 或者 WinRAR 的高版本重试这类工具对 zip 扩展特性的兼容性比系统自带的好。我自己处理过的一个案例是用系统自带的解压一直提示伪加密换 7-Zip 之后正常解出。解压之后最好直接核对一下文件数防止中途缺文件。5.2 Session 创建失败提示算子不存在或 provider 不可用现象ort.InferenceSession创建时抛异常报错信息里出现「Unsupported Operator」或「Execution Provider not found」。原因ONNX 模型的 opset 版本太高当前 onnxruntime 版本不支持或者装的是 CPU 版 onnxruntime但代码里强行指定了 CUDAExecutionProvider。解决优先检查 onnxruntime 版本。如果模型是 opset 18 导出的onnxruntime 至少要 1.16 以上。其次检查 provider 指定逻辑不要写死单一 provider用列表同时声明 CUDA 和 CPU 做回退。这类问题在换机器部署时最容易冒出来源头大都是导出时 opset 设得过于激进我现在的习惯是 opset 统一设 17兼容性最好。5.3 推理结果全是黑图或白图mask 输出异常现象代码正常跑完但输出的 mask 是全黑或者全白看起来像是分割失败但没有任何报错。原因两种典型情况。第一种是输入图像没有做归一化FP32 的原始像素范围是 0 到 255而模型训练时用的是 0 到 1 的归一化值特征分布完全错位。第二种是 prompt 坐标体系不对SAM2 要求 prompt 坐标是输入图 resize 后的绝对坐标如果直接传原图坐标prompt 点落在完全不同的位置上分割结果自然会乱掉。解决检查预处理环节确认图像在送入 encoder 前已经缩放到 1024×1024 并做了归一化prompt 坐标要做同比例的换算。这个环节是黑匣子出错时不报错、不给提示只能靠输出结果反推所以我把 preprocess 和坐标换算单独抽出来做了单元测试。5.4 量化后 mask 边缘像锯齿精度掉了但不报错现象int8 模型跑出来的 mask 形状基本对但边缘全是锯齿小目标直接消失。原因encoder 量化粗粒度导致的特征精度损失。calibration 数据集太少或者分布偏统计的激活值范围不准导致量化误差集中在细节特征上注意力算子的量化损失也会放大边缘噪声。解决回到量化环节增加校准数据到 100 张以上检查 per_channel 是否打开如果还不行就做混合精度把注意力算子的量化排除掉。量化效果看 mIoU 而不是肉眼看图这是血泪经验——有时候肉眼看着差不多但实际上已经掉了三个点。5.5 内存溢出批量推理时图像一次性全载入现象批量推理时提示 OOM 或 MemoryError尤其是带 GPU 的环境更容易触发。原因onnxruntime 的 session 会缓存部分中间张量在多 batch 和大分辨率下显存占用会叠加。特别是 encoder 的输出特征图是四维浮点张量批量一大就超出了显存限额。解决控制推理 batch 大小不要一次性塞几十张图。常见做法是每次喂 4 张或 8 张循环处理。还有一个技巧是复用 session 的run时尽量用同一个输入缓冲区减少张量重新分配的频率。6. 批量分割与结果验证把脚本改造成可落地的推理管线脚本本身跑通之后离生产落地还差一步批量处理多张图、保存结果、量化质量验证。这一节我给出一个可以直接用的批量推理改造版也是我把这套脚本接到实际项目里时做的改动。批量推理的核心是对每张图调用一轮「预处理 → 编码器 → 解码器 → 后处理」把 mask 存成 PNG 文件。prompt 部分我固定用中心点作为前景提示适合「对每张图的主体做分割」这类场景如果你要指定框或多个点改成从表格或 JSON 读取即可import cv2 import glob import numpy as np import onnxruntime as ort # 推理 session 全局复用避免反复创建 enc_session ort.InferenceSession(encoder_int8.onnx, providers[CPUExecutionProvider]) dec_session ort.InferenceSession(decoder.onnx, providers[CPUExecutionProvider]) def infer_one(image_path, out_path, pointNone): img cv2.imread(image_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 预处理缩放、归一化、转 CHW img_resized cv2.resize(img, (1024, 1024)) x img_resized.astype(np.float32) / 255.0 x np.transpose(x, (2, 0, 1))[None] # 变成 1x3x1024x1024 feats enc_session.run([image_features], {images: x})[0] # prompt 默认中心点标签为前景 h, w 1024, 1024 if point is None: point_coords np.array([[[w // 2, h // 2]]], dtypenp.float32) else: point_coords np.array([[point]], dtypenp.float32) point_labels np.ones((1, 1), dtypenp.int64) box np.zeros((1, 4), dtypenp.float32) # 空框不用框提示 masks, ious dec_session.run( [masks, iou_predictions], { image_features: feats, point_coords: point_coords, point_labels: point_labels, box: box } ) # SAM2 输出多个候选 mask取 IoU 最高的那个 best np.argmax(ious[0]) mask masks[0][best] mask (mask 0.0).astype(np.uint8) * 255 cv2.imwrite(out_path, mask) for img_path in glob.glob(images/*.jpg): out_path img_path.replace(images, outputs) infer_one(img_path, out_path, pointNone)这段代码里有三个细节值得说明。第一encoder 和 decoder 两个 session 各自复用不要在循环里反复创建否则光是初始化开销就能让整体耗时翻倍。第二masks的输出形状是(batch, num_masks, H, W)iou_predictions是(batch, num_masks)取argmax选最高置信度的 mask 是 SAM 系列的标准后处理方式。第三mask 阈值用了 0.0因为 SAM2 的 mask 输出已经带 sigmoid 激活大于 0.5 会更严格0.0 能保留更多边缘细节高强度场景可以自己调。批量处理跑通之后我还做了一次全量验证对 200 张测试图分别用 FP32 模型和 int8 模型跑分割计算两版 mask 的 mIoU。那次量化我一开始只用了 10 张校准图结果 mIoU 掉到 0.90边缘大面积破损后来把校准图加到 120 张、打开 per-channelmIoU 回到了 0.96边缘细节基本恢复。从那以后我每次量化都强制走一遍「先校准、再验证、最后看 mIoU」的流程不再凭感觉判断量化好坏。这套流程看起来多花半小时但能省掉后面部署现场排查问题的一整天。希望帮到你。本文还有配套的精品资源点击获取