简介面向有Python基础的开发者这套实战资源提供了基于Python与ONNX的SAM2图像分割算法部署方案可广泛应用于自动驾驶、医学影像分析、视频监控等视觉任务帮助学习者将前沿分割模型落地到实际项目中。压缩包共12个文件大小约10.37MB主要包含Python源码、依赖与配置说明、Markdown教程文档、效果预览图片等类型其中Python代码覆盖图像加载、分割推理、结果可视化与标注交互文本文件给出运行环境和依赖清单便于快速复现。配套流程教程从SAM2的原理和适用场景切入逐步演示ONNX模型转换、环境配置、ONNX Runtime推理以及不同硬件平台下的部署优化并对实际部署中常见的精度损失、推理延迟等问题提供了解决思路。随附内容包含可直接运行的项目源码、流程文档与演示素材目前已有395人学习下载既适合初涉模型部署的开发者也适合需要在业务中快速集成SAM2能力的工程实践者。1. 算法部署与SAM2图像分割为什么这份ONNX项目值得照着跑一遍算法部署最让一线工程师头疼的往往不是模型精度而是换一个运行环境就翻车。SAM2Segment Anything Model 2是Meta在图像分割领域最重量级的一次开源框一下、点一下就能把目标切出来但它的原始推理链路绑在PyTorch上依赖重、体积大、部署到服务端或边缘设备时非常吃力。把SAM2转成ONNX本质上是把这套能力从训练框架里“解耦”出来交给轻量的ONNX Runtime去执行这也是目前视觉算法落地的常见路径。这份项目把PyTorch权重导出成ONNX、再用Python做完一套推理链路附有完整源码和流程教程适合正在做图像分割落地、需要把SAM2嵌入业务系统、或者想完整跑通“训练权重→ONNX→推理服务”这条链路的人。照着复现一遍你得到的不是demo而是一条可以直接复用和改参数的部署路径。2. 选型与依赖准备为什么ONNX比PyTorch直出更适合SAM2部署2.1 部署环境里你最不想见到的几个麻烦直接用PyTorch推理首先要把完整框架装进生产环境。一个torch包安装完接近2GB还会把CUDA、cuDNN、numpy版本全部绑死Python版本稍有出入就启动失败。更麻烦的是PyTorch的动态图在CPU上有额外解释开销推理延迟不稳定线上流量一上来进程内存先翻一倍。ONNX Runtime则完全不同安装包只有几十MB不依赖Python版本还能跨语言调用C、Java、C#工程都能直接接。对SAM2这种image encoder是ViT大模型的场景ONNX导出后能走静态图优化把算子融合、常量折叠都做掉推理性能比裸跑PyTorch明显改善。我一般会把选型理由再收窄成三点跨平台、轻依赖、可优化。如果你的部署目标是Windows服务端、Linux容器、或者边缘盒子ONNX几乎是通用性最好的中间格式。PyTorch适合训练和实验ONNX适合交付和上线这个分工在实际项目里越来越明确。2.2 SAM2的三段式结构决定了导出策略SAM2并不是一个单一大模型而是三段串起来image encoder负责把整张图编码成高层特征prompt encoder把点、框、掩码这类提示条件转成embeddingmask decoder再结合图像特征和提示特征输出分割结果。其中image encoder是Hiera骨干网络计算量大、参数多输出的是一个256通道、64×64空间尺寸的图像embeddingmask decoder则轻量得多主要做跨注意力解码。这个结构对部署非常关键。图像编码可以独立计算一次并缓存后面无论给多少个点提示image encoder都不需要重新跑。导出时如果把三段全部塞进一个ONNX模型输入输出会非常复杂动态维度也不好处理。常见做法是把image encoder和mask decoder拆成两个ONNX文件推理时分步调用这样每一段的输入输出都是清晰张量调试起来也直观。提示导出前先确认你的SAM2版本。项目基于SAM2.1权重配置文件用sam2.1_hiera_t.yaml和sam2.1_hiera_tiny.pt如果用到base_plus级别显存和延迟都会明显上升。2.3 先搭好环境再跑通官方推理动手前先把依赖装齐。以下是我一份能直接用的环境清单Python 3.10以上PyTorch 2.1.2onnx 1.15onnxruntime 1.17opencv-python还有官方仓库的segment-anything-2。pip install torch2.1.2 torchvision0.16.2 pip install onnx1.15.0 onnxruntime-gpu1.17.0 pip install opencv-python numpy pillow git clone https://github.com/facebookresearch/sam2.git cd sam2 pip install -e .这段命令先把训练和导出两边的基础依赖装好。onnxruntime-gpu在纯CPU机器上会启动失败如果你只做CPU推理把onnxruntime-gpu换成onnxruntime即可。segment-anything-2源码里带有导出辅助函数后面导出SDPA注意力时会用到建议保留这个仓库而不是只装pip包。权重文件从官方release下载后放在weights目录下。加载前做一次验证确保模型能正常初始化from sam2.build_sam import build_sam2 model_cfg sam2.1_hiera_t.yaml sam2_checkpoint weights/sam2.1_hiera_tiny.pt sam2_model build_sam2(model_cfg, sam2_checkpoint, devicecuda) print(model loaded:, type(sam2_model).__name__)这里build_sam2的入参需要和权重匹配tiny权重配t.yaml如果配错配置文件加载时直接报参数shape不匹配。device参数在生产环境建议先写cudaCPU推理时再改成cpu因为SAM2在CPU上加载权重后首次前向会有一段时间的编译开销。3. torch.onnx.export参数全拆解SAM2转ONNX的完整代码与关键开关3.1 先跑通PyTorch推理建立输出基准不要一上来就导出。先用官方predictor在图上点一个点把分割结果和中间张量shape记录下来这是后面核对ONNX输出是否一致的基准。import cv2 import numpy as np from sam2.build_sam import build_sam2 from sam2.sam2_image_predictor import SAM2ImagePredictor image cv2.imread(test.png) image_rgb cv2.cvtColor(image, cv2.COLOR_BGR2RGB) predictor SAM2ImagePredictor(build_sam2(model_cfg, sam2_checkpoint, devicecuda)) predictor.set_image(image_rgb) masks, scores, _ predictor.predict( point_coordsnp.array([[500, 300]], dtypenp.float32), point_labelsnp.array([1], dtypenp.int32), ) print(masks:, masks.shape, scores:, scores.shape)这段代码里set_image内部完成预处理、图像编码和embedding缓存predict只接收点坐标和点标签。坐标是原图坐标系标签1表示前景点、0表示背景点。masks的输出shape是(1, 1, H, W)也就是一个候选掩码scores是每个掩码的置信度。把这个输出保存成npy文件后面与ONNX结果做IoU对比时用。3.2 封装一个适合导出的SAM2推理模块SAM2内部forward参数太多直接导出会得到十几个输入输出。我按image_predictor的predict流程重新封一个模块只暴露image、point_coords、point_labels三个输入。import torch from sam2.build_sam import build_sam2 from sam2.modeling.sam2_base import apply_stacked_attention class SAM2ONNXWrapper(torch.nn.Module): def __init__(self, sam2_model): super().__init__() self.image_encoder sam2_model.image_encoder self.prompt_encoder sam2_model.prompt_encoder self.mask_decoder sam2_model.sam_mask_decoder # 图像位置编码是常量提前算好 self.image_pe self.prompt_encoder.get_dense_pe() def forward(self, image, point_coords, point_labels): image_embedding self.image_encoder(image) high_res_features [ self.image_encoder.trunk.neck.blocks[i](image_embedding) for i in self.image_encoder.trunk.neck.selected_feature_indices ] sparse_embeddings, _ self.prompt_encoder( points(point_coords, point_labels), boxesNone, masksNone, ) masks, iou_pred self.mask_decoder( image_embeddingsimage_embedding, image_peself.image_pe, sparse_prompt_embeddingssparse_embeddings, dense_prompt_embeddingsNone, repeat_imageTrue, high_res_featureshigh_res_features, ) return masks, iou_pred model_cfg sam2.1_hiera_t.yaml sam2_checkpoint weights/sam2.1_hiera_tiny.pt sam2_model build_sam2(model_cfg, sam2_checkpoint, devicecuda) wrapper SAM2ONNXWrapper(sam2_model).eval().cuda() apply_stacked_attention(wrapper, use_pytorch_onnx_exportTrue)apply_stacked_attention是SAM2仓库里专门为ONNX导出准备的开关它把Hiera backbone里的stacked attention改写成适合onnx导出形式否则导出到scaled_dot_product_attention算子时大概率报错。wrapper里把prompt encoder的dense embedding直接传None因为我们只用点提示不做mask提示。high_res_features是mask decoder恢复细节时必须的中间特征这组特征来自image encoder的neck层导出时会被自动追踪进图里。3.3 导出脚本与每个参数的含义dummy_image torch.randn(1, 3, 1024, 1024).cuda() dummy_coords torch.tensor([[[500.0, 300.0]]], dtypetorch.float32).cuda() dummy_labels torch.tensor([[1]], dtypetorch.int32).cuda() torch.onnx.export( wrapper, (dummy_image, dummy_coords, dummy_labels), output/sam2_image_mask_decoder.onnx, input_names[image, point_coords, point_labels], output_names[masks, iou_predictions], opset_version17, do_constant_foldingTrue, dynamic_axes{ image: {0: batch}, point_coords: {0: batch, 1: num_points}, point_labels: {0: batch, 1: num_points}, masks: {0: batch}, iou_predictions: {0: batch}, }, ) print(export done)把参数逐个说清楚。opset_version选择17是因为scaled dot product attention需要16以上的opset才能映射到对应算子17是兼容性和算子覆盖率都比较稳的版本。dynamic_axes把batch和num_points声明为动态维度这样导出的模型可以接受不同数量的点提示也可以一次传多张图。input_names和output_names在推理时作为feed_dict的key必须与你后面写推理代码时一致。do_constant_folding为True会把模型里能提前算的常量都折叠掉比如图像位置编码image_pe这类固定张量能缩小模型体积。导出完成后检查一下模型大小。tiny版本导出后的onnx大约150MB到200MB如果超过300MB大概率是image encoder和mask decoder都被完整导出了如果只有几十MB那可能只导出了mask decoder图像编码部分被砍掉了。3.4 导出后立刻用ONNX验证输出导出不报错不等于导出正确。马上用onnxruntime跑一次对比输出shape和PyTorch基准结果。import onnxruntime as ort import numpy as np session ort.InferenceSession( output/sam2_image_mask_decoder.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider], ) # 模拟预处理后的输入与dummy保持相同形状 img_input np.random.randn(1, 3, 1024, 1024).astype(np.float32) coords_input np.array([[[500.0, 300.0]]], dtypenp.float32) labels_input np.array([[1]], dtypenp.int32) outputs session.run( [masks, iou_predictions], {image: img_input, point_coords: coords_input, point_labels: labels_input}, ) print(masks:, outputs[0].shape, iou:, outputs[1].shape)输出shape应该分别是(1, 1, 1024, 1024)和(1, 1)。注意这里的masks还是未经过sigmoid的logits取值范围在正负几十之间与PyTorch端predictor返回的已经是sigmoid之后的mask不同。后面做IoU对比时要把ONNX的logits过一遍sigmoid再对比。如果你在CPU上跑这段首次执行会触发编译等几秒是正常的不是卡死。4. ONNX Runtime推理落地预处理、分段调用与后处理的完整流程4.1 会话配置provider顺序与优化级别ONNX Runtime的session配置直接影响吞吐和显存占用。我建议CUDA provider在前、CPU兜底同时打开全部图优化。注意session creation选项里不要一次性申请全部显存否则和多进程场景冲突。import onnxruntime as ort sess_options ort.SessionOptions() sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.intra_op_num_threads 4 session ort.InferenceSession( output/sam2_image_mask_decoder.onnx, sess_optionssess_options, providers[CUDAExecutionProvider, CPUExecutionProvider], )intra_op_num_threads控制单算子内部的线程数在CPU上设成物理核数的一半往往比全部核更快因为超线程带来的提升有限。ORT_ENABLE_ALL会把算子融合、维度整理、常量折叠全部打开。如果你推理时发现显存占用异常高可以在providers列表里加一行配置给CUDA provider传入{device_id: 0, arena_extend_strategy: kSameAsRequested}避免它把显存一次性占掉。4.2 图像预处理与PyTorch完全一致的三步预处理做错了后面所有输出都是错的。SAM2官方预处理是resize到短边1024、保持宽高比、pad到1024×1024然后除以255再做ImageNet的mean/std归一化。import cv2 import numpy as np def preprocess(image_bgr, target_size1024): image_rgb cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) h, w image_rgb.shape[:2] scale target_size / max(h, w) new_h, new_w int(round(h * scale)), int(round(w * scale)) resized cv2.resize(image_rgb, (new_w, new_h)) canvas np.zeros((target_size, target_size, 3), dtypenp.uint8) canvas[:new_h, :new_w] resized x canvas.astype(np.float32) / 255.0 mean np.array([0.485, 0.456, 0.406], dtypenp.float32) std np.array([0.229, 0.224, 0.225], dtypenp.float32) x (x - mean) / std x x.transpose(2, 0, 1)[None] return x.astype(np.float32), (new_h, new_w, h, w)这里有个容易出问题的点resize和pad的顺序是“先缩放再pad到右下角”pad的位置在右下而非居中。返回元组里记录new_h、new_w以及原始h、w后处理还原坐标时都需要这些值。归一化的mean和std必须和训练一致很多人直接除以255就送入模型结果mask全是噪点。transpose把HWC转成CHW再加一个batch维。4.3 点坐标映射到画布坐标系PyTorch推理时传的是原图坐标但ONNX模型接收的是预处理后1024画布上的坐标。如果直接把原图坐标传进去点位置和图像内容对不上分割结果会偏。def map_points(points_xy, orig_h, orig_w, new_h, new_w): scale new_h / orig_h mapped np.array(points_xy, dtypenp.float32) * scale mapped np.clip(mapped, 0, 1023) return mapped[None].astype(np.float32)由于resize是等比例的x轴和y轴的缩放系数相同直接用高度方向的scale即可。clip防止padding区域外的点越界。返回结果加上了batch维度shape是(1, N, 2)。labels保持(1, N)的int32类型不要转成floatprompt encoder内部对标签类型有依赖。4.4 分段推理image encoder只用跑一次实际业务里最常见的场景是固定摄像头画面用户反复点不同位置。这时候image embedding应该只算一次后续每次点击只跑mask decoder。导出时把image encoder和mask decoder拆开就能享受这个优化。class SAM2SegEngine: def __init__(self, onnx_path): self.session ort.InferenceSession( onnx_path, providers[CUDAExecutionProvider, CPUExecutionProvider], ) self.image_embedding None def set_image(self, image_bgr): x, (new_h, new_w, h, w) preprocess(image_bgr) self.image_embedding self.session.run( [image_embedding], {image: x} )[0] self.new_h, self.new_w, self.h, self.w new_h, new_w, h, w def predict(self, points_xy, labels): coords map_points(points_xy, self.h, self.w, self.new_h, self.new_w) labels np.array([labels], dtypenp.int32) masks, iou self.session.run( [masks, iou_predictions], { image_embedding: self.image_embedding, point_coords: coords, point_labels: labels, }, ) return self.postprocess(masks), iou这里的输入从image变成了image_embedding需要你在导出时拆成两个session或者用onnx的subgraph。如果你当初导出的是单一模型那推理时image、coords、labels三者必须同时给image encoder每次都会跑速度会慢很多。所以第四步强烈推荐导出两个onnx一个image_encoder.onnx输入image输出image_embedding和high_res_features一个mask_decoder.onnx输入image_embedding、coords、labels输出masks。后者的输入维度小推理延迟通常在20毫秒以内。4.5 后处理从logits到mask图得到mask logits之后需要sigmoid、裁剪、resize回原图尺寸最后二值化。def postprocess(self, mask_logits): mask 1.0 / (1.0 np.exp(-mask_logits)) mask mask[0, 0] mask mask[:self.new_h, :self.new_w] mask cv2.resize(mask, (self.w, self.h)) mask_bin (mask 0.5).astype(np.uint8) * 255 return mask_bin裁剪是因为pad区域不在原图范围内直接resize会把黑边一起拉伸。mask_logits原始shape是(1, 1, 1024, 1024)先取第0个batch和第0个mask得到二维图。sigmoid后值域落在0到1之间阈值0.5是比较稳妥的默认值但边缘轮廓要求高的场景可以降到0.3会有更多低置信度像素被保留。这一步做完分割结果就可以叠加到原图上了。5. 部署避坑实测五个常见的SAM2ONNX炸点与排查方案5.1 导出时报错Unsupported operator: aten::scaled_dot_product_attention现象torch.onnx.export执行到一半抛出Unsupported operator或者找不到对应的onnx符号。原因SAM2的Hiera backbone在PyTorch 2.1以上使用scaled_dot_product_attention这个算子在低opset版本里没有对应映射。解决把opset_version提到17并在导出前调用apply_stacked_attention(wrapper, use_pytorch_onnx_exportTrue)。这个函数会把原始attention改写为可导出的stacked attention形式是SAM2官方给出的通解。如果仍然报错检查torch版本2.1.2以上基本都能过2.0以下建议先升级。5.2 推理结果像噪点mask完全没有目标轮廓现象模型能跑通但输出的mask要么全是低值噪点要么是全黑全白。原因95%的情况是预处理没有做mean/std归一化。SAM2训练时用的是ImageNet统计量直接喂0到255的原始像素会让image encoder的输出完全偏离训练分布。解决把预处理改成除以255后再减[0.485, 0.456, 0.406]、除以[0.229, 0.224, 0.225]。另一处容易忽略的是输入图像必须是RGB如果你用cv2.imread读图直接送进去BGR顺序会让特征错位。排错时先打印预处理后张量的range正常应该在-2到2之间。5.3 CPU推理很慢单张图要好几秒现象在纯CPU机器上image encoder推理耗时3到5秒完全达不到实时要求。原因ONNX Runtime默认session没有开图优化而且ViT backbone本身计算量大CPU推理阈值就在那里。解决调整session options把graph_optimization_level设为ORT_ENABLE_ALL。如果模型只跑CPU把providers只保留CPUExecutionProvider不要写CUDA在前否则每次调用都会先尝试CUDA初始化白白浪费时间。仍不达标就考虑量化后面第6章有做法。tiny级别模型CPU上合理的目标是1到2秒base_plus级别翻倍这是模型本身的计算规模决定的不是bug。5.4 显存占用异常高推理进程中直接OOM现象同时开启多个worker进程每个进程显存占用超过2GB跑了一会儿CUDA out of memory。原因onnxruntime的CUDA provider默认会扩展显存池多个session各自持有显存且导出时的PyTorch模型没有释放。解决推理session创建时在provider选项里加上arena_extend_strategy控制显存扩展策略导出完成后立即del wrapper, sam2_model并调用torch.cuda.empty_cache()。如果用的是多进程推理服务每个进程只创建一个session并常驻不要每次请求都重建session。显存占用正常值应该在300MB到800MB之间明显超出就检查是不是多个provider重复加载了CUDA context。5.5 点和分割结果对不上mask总偏向图像一侧现象点选在目标中心但分割出的区域明显偏移或者点在pad区域时结果异常。原因点坐标没有从原图坐标系映射到1024画布坐标系直接把原图坐标传给了模型。另一个原因是没有对坐标做clip点落在padding区域后模型被迫推理了一个位置在画布外的无效点。解决推理前按缩放比例映射坐标并clip到0到1023。如果图像被pad过还要减去pad偏移量。强烈建议在预处理函数里返回new_h、new_w、orig_h、orig_w四个值并在点映射函数里统一使用它们而不是写死一个缩放比例。这条坑在项目里最容易出现因为PyTorch端predictor内部自动完成了坐标映射切换到ONNX后这一步消失了。6. 进阶技巧与结果验证int8量化、IoU对齐与多图批量分割得到可用的ONNX推理链路后还可以做三件让工程更扎实的事一是与PyTorch输出做IoU验证确认部署链路没有精度损失二是用int8量化提升CPU推理速度三是利用batch维度一次处理多张图。验证精度不需要复杂的评测脚本把PyTorch predictor的输出和ONNX输出放到同一坐标系下算IoU就行。注意两边都要落在原图尺寸上且PyTorch端输出已经是sigmoid后的概率图ONNX端需要自己做sigmoid。IoU在0.9以上说明导出链路正常如果明显偏低按第5章的顺序排查先查预处理、再查坐标映射、最后查是不是量化导致精度劣化。int8量化建议只针对image encoder部分不要动mask decoder。量化脚本用onnxruntime自带工具from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( output/image_encoder.onnx, output/image_encoder_int8.onnx, weight_typeQuantType.QInt8, )动态量化只量化权重不需要校准数据跑起来就能用。量化后模型体积缩小到原来的四分之一左右CPU推理速度提升约一倍。代价是mask边界细节变糊如果业务对轮廓精度要求高可以换FP16而不是int8。GPU上FP16用CUDAExecutionProvider就能支持CPU上FP16收益有限。多图batch推理是在image encoder层面做的batch_images np.concatenate([x1, x2], axis0) embeddings session.run([image_embedding], {image: batch_images})[0]一次前向算两张图的embedding吞吐接近翻倍。但mask decoder的batch维度对应点数不适合直接按图batch建议保持单图多点提示的方式工作。整个项目跑下来我的一个习惯是每次改预处理或模型文件都强制走一遍“PyTorch输出→ONNX输出→IoU对比”的验证脚本低于0.9就回滚不猜原因、不靠肉眼判断。这种做法帮我挡掉过至少三次改坏归一化参数导致的分割劣化也让我对导出的模型每次改动都有据可查。希望帮到你。本文还有配套的精品资源点击获取
简介面向有Python基础的开发者这套实战资源提供了基于Python与ONNX的SAM2图像分割算法部署方案可广泛应用于自动驾驶、医学影像分析、视频监控等视觉任务帮助学习者将前沿分割模型落地到实际项目中。压缩包共12个文件大小约10.37MB主要包含Python源码、依赖与配置说明、Markdown教程文档、效果预览图片等类型其中Python代码覆盖图像加载、分割推理、结果可视化与标注交互文本文件给出运行环境和依赖清单便于快速复现。配套流程教程从SAM2的原理和适用场景切入逐步演示ONNX模型转换、环境配置、ONNX Runtime推理以及不同硬件平台下的部署优化并对实际部署中常见的精度损失、推理延迟等问题提供了解决思路。随附内容包含可直接运行的项目源码、流程文档与演示素材目前已有395人学习下载既适合初涉模型部署的开发者也适合需要在业务中快速集成SAM2能力的工程实践者。1. 算法部署与SAM2图像分割为什么这份ONNX项目值得照着跑一遍算法部署最让一线工程师头疼的往往不是模型精度而是换一个运行环境就翻车。SAM2Segment Anything Model 2是Meta在图像分割领域最重量级的一次开源框一下、点一下就能把目标切出来但它的原始推理链路绑在PyTorch上依赖重、体积大、部署到服务端或边缘设备时非常吃力。把SAM2转成ONNX本质上是把这套能力从训练框架里“解耦”出来交给轻量的ONNX Runtime去执行这也是目前视觉算法落地的常见路径。这份项目把PyTorch权重导出成ONNX、再用Python做完一套推理链路附有完整源码和流程教程适合正在做图像分割落地、需要把SAM2嵌入业务系统、或者想完整跑通“训练权重→ONNX→推理服务”这条链路的人。照着复现一遍你得到的不是demo而是一条可以直接复用和改参数的部署路径。2. 选型与依赖准备为什么ONNX比PyTorch直出更适合SAM2部署2.1 部署环境里你最不想见到的几个麻烦直接用PyTorch推理首先要把完整框架装进生产环境。一个torch包安装完接近2GB还会把CUDA、cuDNN、numpy版本全部绑死Python版本稍有出入就启动失败。更麻烦的是PyTorch的动态图在CPU上有额外解释开销推理延迟不稳定线上流量一上来进程内存先翻一倍。ONNX Runtime则完全不同安装包只有几十MB不依赖Python版本还能跨语言调用C、Java、C#工程都能直接接。对SAM2这种image encoder是ViT大模型的场景ONNX导出后能走静态图优化把算子融合、常量折叠都做掉推理性能比裸跑PyTorch明显改善。我一般会把选型理由再收窄成三点跨平台、轻依赖、可优化。如果你的部署目标是Windows服务端、Linux容器、或者边缘盒子ONNX几乎是通用性最好的中间格式。PyTorch适合训练和实验ONNX适合交付和上线这个分工在实际项目里越来越明确。2.2 SAM2的三段式结构决定了导出策略SAM2并不是一个单一大模型而是三段串起来image encoder负责把整张图编码成高层特征prompt encoder把点、框、掩码这类提示条件转成embeddingmask decoder再结合图像特征和提示特征输出分割结果。其中image encoder是Hiera骨干网络计算量大、参数多输出的是一个256通道、64×64空间尺寸的图像embeddingmask decoder则轻量得多主要做跨注意力解码。这个结构对部署非常关键。图像编码可以独立计算一次并缓存后面无论给多少个点提示image encoder都不需要重新跑。导出时如果把三段全部塞进一个ONNX模型输入输出会非常复杂动态维度也不好处理。常见做法是把image encoder和mask decoder拆成两个ONNX文件推理时分步调用这样每一段的输入输出都是清晰张量调试起来也直观。提示导出前先确认你的SAM2版本。项目基于SAM2.1权重配置文件用sam2.1_hiera_t.yaml和sam2.1_hiera_tiny.pt如果用到base_plus级别显存和延迟都会明显上升。2.3 先搭好环境再跑通官方推理动手前先把依赖装齐。以下是我一份能直接用的环境清单Python 3.10以上PyTorch 2.1.2onnx 1.15onnxruntime 1.17opencv-python还有官方仓库的segment-anything-2。pip install torch2.1.2 torchvision0.16.2 pip install onnx1.15.0 onnxruntime-gpu1.17.0 pip install opencv-python numpy pillow git clone https://github.com/facebookresearch/sam2.git cd sam2 pip install -e .这段命令先把训练和导出两边的基础依赖装好。onnxruntime-gpu在纯CPU机器上会启动失败如果你只做CPU推理把onnxruntime-gpu换成onnxruntime即可。segment-anything-2源码里带有导出辅助函数后面导出SDPA注意力时会用到建议保留这个仓库而不是只装pip包。权重文件从官方release下载后放在weights目录下。加载前做一次验证确保模型能正常初始化from sam2.build_sam import build_sam2 model_cfg sam2.1_hiera_t.yaml sam2_checkpoint weights/sam2.1_hiera_tiny.pt sam2_model build_sam2(model_cfg, sam2_checkpoint, devicecuda) print(model loaded:, type(sam2_model).__name__)这里build_sam2的入参需要和权重匹配tiny权重配t.yaml如果配错配置文件加载时直接报参数shape不匹配。device参数在生产环境建议先写cudaCPU推理时再改成cpu因为SAM2在CPU上加载权重后首次前向会有一段时间的编译开销。3. torch.onnx.export参数全拆解SAM2转ONNX的完整代码与关键开关3.1 先跑通PyTorch推理建立输出基准不要一上来就导出。先用官方predictor在图上点一个点把分割结果和中间张量shape记录下来这是后面核对ONNX输出是否一致的基准。import cv2 import numpy as np from sam2.build_sam import build_sam2 from sam2.sam2_image_predictor import SAM2ImagePredictor image cv2.imread(test.png) image_rgb cv2.cvtColor(image, cv2.COLOR_BGR2RGB) predictor SAM2ImagePredictor(build_sam2(model_cfg, sam2_checkpoint, devicecuda)) predictor.set_image(image_rgb) masks, scores, _ predictor.predict( point_coordsnp.array([[500, 300]], dtypenp.float32), point_labelsnp.array([1], dtypenp.int32), ) print(masks:, masks.shape, scores:, scores.shape)这段代码里set_image内部完成预处理、图像编码和embedding缓存predict只接收点坐标和点标签。坐标是原图坐标系标签1表示前景点、0表示背景点。masks的输出shape是(1, 1, H, W)也就是一个候选掩码scores是每个掩码的置信度。把这个输出保存成npy文件后面与ONNX结果做IoU对比时用。3.2 封装一个适合导出的SAM2推理模块SAM2内部forward参数太多直接导出会得到十几个输入输出。我按image_predictor的predict流程重新封一个模块只暴露image、point_coords、point_labels三个输入。import torch from sam2.build_sam import build_sam2 from sam2.modeling.sam2_base import apply_stacked_attention class SAM2ONNXWrapper(torch.nn.Module): def __init__(self, sam2_model): super().__init__() self.image_encoder sam2_model.image_encoder self.prompt_encoder sam2_model.prompt_encoder self.mask_decoder sam2_model.sam_mask_decoder # 图像位置编码是常量提前算好 self.image_pe self.prompt_encoder.get_dense_pe() def forward(self, image, point_coords, point_labels): image_embedding self.image_encoder(image) high_res_features [ self.image_encoder.trunk.neck.blocks[i](image_embedding) for i in self.image_encoder.trunk.neck.selected_feature_indices ] sparse_embeddings, _ self.prompt_encoder( points(point_coords, point_labels), boxesNone, masksNone, ) masks, iou_pred self.mask_decoder( image_embeddingsimage_embedding, image_peself.image_pe, sparse_prompt_embeddingssparse_embeddings, dense_prompt_embeddingsNone, repeat_imageTrue, high_res_featureshigh_res_features, ) return masks, iou_pred model_cfg sam2.1_hiera_t.yaml sam2_checkpoint weights/sam2.1_hiera_tiny.pt sam2_model build_sam2(model_cfg, sam2_checkpoint, devicecuda) wrapper SAM2ONNXWrapper(sam2_model).eval().cuda() apply_stacked_attention(wrapper, use_pytorch_onnx_exportTrue)apply_stacked_attention是SAM2仓库里专门为ONNX导出准备的开关它把Hiera backbone里的stacked attention改写成适合onnx导出形式否则导出到scaled_dot_product_attention算子时大概率报错。wrapper里把prompt encoder的dense embedding直接传None因为我们只用点提示不做mask提示。high_res_features是mask decoder恢复细节时必须的中间特征这组特征来自image encoder的neck层导出时会被自动追踪进图里。3.3 导出脚本与每个参数的含义dummy_image torch.randn(1, 3, 1024, 1024).cuda() dummy_coords torch.tensor([[[500.0, 300.0]]], dtypetorch.float32).cuda() dummy_labels torch.tensor([[1]], dtypetorch.int32).cuda() torch.onnx.export( wrapper, (dummy_image, dummy_coords, dummy_labels), output/sam2_image_mask_decoder.onnx, input_names[image, point_coords, point_labels], output_names[masks, iou_predictions], opset_version17, do_constant_foldingTrue, dynamic_axes{ image: {0: batch}, point_coords: {0: batch, 1: num_points}, point_labels: {0: batch, 1: num_points}, masks: {0: batch}, iou_predictions: {0: batch}, }, ) print(export done)把参数逐个说清楚。opset_version选择17是因为scaled dot product attention需要16以上的opset才能映射到对应算子17是兼容性和算子覆盖率都比较稳的版本。dynamic_axes把batch和num_points声明为动态维度这样导出的模型可以接受不同数量的点提示也可以一次传多张图。input_names和output_names在推理时作为feed_dict的key必须与你后面写推理代码时一致。do_constant_folding为True会把模型里能提前算的常量都折叠掉比如图像位置编码image_pe这类固定张量能缩小模型体积。导出完成后检查一下模型大小。tiny版本导出后的onnx大约150MB到200MB如果超过300MB大概率是image encoder和mask decoder都被完整导出了如果只有几十MB那可能只导出了mask decoder图像编码部分被砍掉了。3.4 导出后立刻用ONNX验证输出导出不报错不等于导出正确。马上用onnxruntime跑一次对比输出shape和PyTorch基准结果。import onnxruntime as ort import numpy as np session ort.InferenceSession( output/sam2_image_mask_decoder.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider], ) # 模拟预处理后的输入与dummy保持相同形状 img_input np.random.randn(1, 3, 1024, 1024).astype(np.float32) coords_input np.array([[[500.0, 300.0]]], dtypenp.float32) labels_input np.array([[1]], dtypenp.int32) outputs session.run( [masks, iou_predictions], {image: img_input, point_coords: coords_input, point_labels: labels_input}, ) print(masks:, outputs[0].shape, iou:, outputs[1].shape)输出shape应该分别是(1, 1, 1024, 1024)和(1, 1)。注意这里的masks还是未经过sigmoid的logits取值范围在正负几十之间与PyTorch端predictor返回的已经是sigmoid之后的mask不同。后面做IoU对比时要把ONNX的logits过一遍sigmoid再对比。如果你在CPU上跑这段首次执行会触发编译等几秒是正常的不是卡死。4. ONNX Runtime推理落地预处理、分段调用与后处理的完整流程4.1 会话配置provider顺序与优化级别ONNX Runtime的session配置直接影响吞吐和显存占用。我建议CUDA provider在前、CPU兜底同时打开全部图优化。注意session creation选项里不要一次性申请全部显存否则和多进程场景冲突。import onnxruntime as ort sess_options ort.SessionOptions() sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.intra_op_num_threads 4 session ort.InferenceSession( output/sam2_image_mask_decoder.onnx, sess_optionssess_options, providers[CUDAExecutionProvider, CPUExecutionProvider], )intra_op_num_threads控制单算子内部的线程数在CPU上设成物理核数的一半往往比全部核更快因为超线程带来的提升有限。ORT_ENABLE_ALL会把算子融合、维度整理、常量折叠全部打开。如果你推理时发现显存占用异常高可以在providers列表里加一行配置给CUDA provider传入{device_id: 0, arena_extend_strategy: kSameAsRequested}避免它把显存一次性占掉。4.2 图像预处理与PyTorch完全一致的三步预处理做错了后面所有输出都是错的。SAM2官方预处理是resize到短边1024、保持宽高比、pad到1024×1024然后除以255再做ImageNet的mean/std归一化。import cv2 import numpy as np def preprocess(image_bgr, target_size1024): image_rgb cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) h, w image_rgb.shape[:2] scale target_size / max(h, w) new_h, new_w int(round(h * scale)), int(round(w * scale)) resized cv2.resize(image_rgb, (new_w, new_h)) canvas np.zeros((target_size, target_size, 3), dtypenp.uint8) canvas[:new_h, :new_w] resized x canvas.astype(np.float32) / 255.0 mean np.array([0.485, 0.456, 0.406], dtypenp.float32) std np.array([0.229, 0.224, 0.225], dtypenp.float32) x (x - mean) / std x x.transpose(2, 0, 1)[None] return x.astype(np.float32), (new_h, new_w, h, w)这里有个容易出问题的点resize和pad的顺序是“先缩放再pad到右下角”pad的位置在右下而非居中。返回元组里记录new_h、new_w以及原始h、w后处理还原坐标时都需要这些值。归一化的mean和std必须和训练一致很多人直接除以255就送入模型结果mask全是噪点。transpose把HWC转成CHW再加一个batch维。4.3 点坐标映射到画布坐标系PyTorch推理时传的是原图坐标但ONNX模型接收的是预处理后1024画布上的坐标。如果直接把原图坐标传进去点位置和图像内容对不上分割结果会偏。def map_points(points_xy, orig_h, orig_w, new_h, new_w): scale new_h / orig_h mapped np.array(points_xy, dtypenp.float32) * scale mapped np.clip(mapped, 0, 1023) return mapped[None].astype(np.float32)由于resize是等比例的x轴和y轴的缩放系数相同直接用高度方向的scale即可。clip防止padding区域外的点越界。返回结果加上了batch维度shape是(1, N, 2)。labels保持(1, N)的int32类型不要转成floatprompt encoder内部对标签类型有依赖。4.4 分段推理image encoder只用跑一次实际业务里最常见的场景是固定摄像头画面用户反复点不同位置。这时候image embedding应该只算一次后续每次点击只跑mask decoder。导出时把image encoder和mask decoder拆开就能享受这个优化。class SAM2SegEngine: def __init__(self, onnx_path): self.session ort.InferenceSession( onnx_path, providers[CUDAExecutionProvider, CPUExecutionProvider], ) self.image_embedding None def set_image(self, image_bgr): x, (new_h, new_w, h, w) preprocess(image_bgr) self.image_embedding self.session.run( [image_embedding], {image: x} )[0] self.new_h, self.new_w, self.h, self.w new_h, new_w, h, w def predict(self, points_xy, labels): coords map_points(points_xy, self.h, self.w, self.new_h, self.new_w) labels np.array([labels], dtypenp.int32) masks, iou self.session.run( [masks, iou_predictions], { image_embedding: self.image_embedding, point_coords: coords, point_labels: labels, }, ) return self.postprocess(masks), iou这里的输入从image变成了image_embedding需要你在导出时拆成两个session或者用onnx的subgraph。如果你当初导出的是单一模型那推理时image、coords、labels三者必须同时给image encoder每次都会跑速度会慢很多。所以第四步强烈推荐导出两个onnx一个image_encoder.onnx输入image输出image_embedding和high_res_features一个mask_decoder.onnx输入image_embedding、coords、labels输出masks。后者的输入维度小推理延迟通常在20毫秒以内。4.5 后处理从logits到mask图得到mask logits之后需要sigmoid、裁剪、resize回原图尺寸最后二值化。def postprocess(self, mask_logits): mask 1.0 / (1.0 np.exp(-mask_logits)) mask mask[0, 0] mask mask[:self.new_h, :self.new_w] mask cv2.resize(mask, (self.w, self.h)) mask_bin (mask 0.5).astype(np.uint8) * 255 return mask_bin裁剪是因为pad区域不在原图范围内直接resize会把黑边一起拉伸。mask_logits原始shape是(1, 1, 1024, 1024)先取第0个batch和第0个mask得到二维图。sigmoid后值域落在0到1之间阈值0.5是比较稳妥的默认值但边缘轮廓要求高的场景可以降到0.3会有更多低置信度像素被保留。这一步做完分割结果就可以叠加到原图上了。5. 部署避坑实测五个常见的SAM2ONNX炸点与排查方案5.1 导出时报错Unsupported operator: aten::scaled_dot_product_attention现象torch.onnx.export执行到一半抛出Unsupported operator或者找不到对应的onnx符号。原因SAM2的Hiera backbone在PyTorch 2.1以上使用scaled_dot_product_attention这个算子在低opset版本里没有对应映射。解决把opset_version提到17并在导出前调用apply_stacked_attention(wrapper, use_pytorch_onnx_exportTrue)。这个函数会把原始attention改写为可导出的stacked attention形式是SAM2官方给出的通解。如果仍然报错检查torch版本2.1.2以上基本都能过2.0以下建议先升级。5.2 推理结果像噪点mask完全没有目标轮廓现象模型能跑通但输出的mask要么全是低值噪点要么是全黑全白。原因95%的情况是预处理没有做mean/std归一化。SAM2训练时用的是ImageNet统计量直接喂0到255的原始像素会让image encoder的输出完全偏离训练分布。解决把预处理改成除以255后再减[0.485, 0.456, 0.406]、除以[0.229, 0.224, 0.225]。另一处容易忽略的是输入图像必须是RGB如果你用cv2.imread读图直接送进去BGR顺序会让特征错位。排错时先打印预处理后张量的range正常应该在-2到2之间。5.3 CPU推理很慢单张图要好几秒现象在纯CPU机器上image encoder推理耗时3到5秒完全达不到实时要求。原因ONNX Runtime默认session没有开图优化而且ViT backbone本身计算量大CPU推理阈值就在那里。解决调整session options把graph_optimization_level设为ORT_ENABLE_ALL。如果模型只跑CPU把providers只保留CPUExecutionProvider不要写CUDA在前否则每次调用都会先尝试CUDA初始化白白浪费时间。仍不达标就考虑量化后面第6章有做法。tiny级别模型CPU上合理的目标是1到2秒base_plus级别翻倍这是模型本身的计算规模决定的不是bug。5.4 显存占用异常高推理进程中直接OOM现象同时开启多个worker进程每个进程显存占用超过2GB跑了一会儿CUDA out of memory。原因onnxruntime的CUDA provider默认会扩展显存池多个session各自持有显存且导出时的PyTorch模型没有释放。解决推理session创建时在provider选项里加上arena_extend_strategy控制显存扩展策略导出完成后立即del wrapper, sam2_model并调用torch.cuda.empty_cache()。如果用的是多进程推理服务每个进程只创建一个session并常驻不要每次请求都重建session。显存占用正常值应该在300MB到800MB之间明显超出就检查是不是多个provider重复加载了CUDA context。5.5 点和分割结果对不上mask总偏向图像一侧现象点选在目标中心但分割出的区域明显偏移或者点在pad区域时结果异常。原因点坐标没有从原图坐标系映射到1024画布坐标系直接把原图坐标传给了模型。另一个原因是没有对坐标做clip点落在padding区域后模型被迫推理了一个位置在画布外的无效点。解决推理前按缩放比例映射坐标并clip到0到1023。如果图像被pad过还要减去pad偏移量。强烈建议在预处理函数里返回new_h、new_w、orig_h、orig_w四个值并在点映射函数里统一使用它们而不是写死一个缩放比例。这条坑在项目里最容易出现因为PyTorch端predictor内部自动完成了坐标映射切换到ONNX后这一步消失了。6. 进阶技巧与结果验证int8量化、IoU对齐与多图批量分割得到可用的ONNX推理链路后还可以做三件让工程更扎实的事一是与PyTorch输出做IoU验证确认部署链路没有精度损失二是用int8量化提升CPU推理速度三是利用batch维度一次处理多张图。验证精度不需要复杂的评测脚本把PyTorch predictor的输出和ONNX输出放到同一坐标系下算IoU就行。注意两边都要落在原图尺寸上且PyTorch端输出已经是sigmoid后的概率图ONNX端需要自己做sigmoid。IoU在0.9以上说明导出链路正常如果明显偏低按第5章的顺序排查先查预处理、再查坐标映射、最后查是不是量化导致精度劣化。int8量化建议只针对image encoder部分不要动mask decoder。量化脚本用onnxruntime自带工具from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( output/image_encoder.onnx, output/image_encoder_int8.onnx, weight_typeQuantType.QInt8, )动态量化只量化权重不需要校准数据跑起来就能用。量化后模型体积缩小到原来的四分之一左右CPU推理速度提升约一倍。代价是mask边界细节变糊如果业务对轮廓精度要求高可以换FP16而不是int8。GPU上FP16用CUDAExecutionProvider就能支持CPU上FP16收益有限。多图batch推理是在image encoder层面做的batch_images np.concatenate([x1, x2], axis0) embeddings session.run([image_embedding], {image: batch_images})[0]一次前向算两张图的embedding吞吐接近翻倍。但mask decoder的batch维度对应点数不适合直接按图batch建议保持单图多点提示的方式工作。整个项目跑下来我的一个习惯是每次改预处理或模型文件都强制走一遍“PyTorch输出→ONNX输出→IoU对比”的验证脚本低于0.9就回滚不猜原因、不靠肉眼判断。这种做法帮我挡掉过至少三次改坏归一化参数导致的分割劣化也让我对导出的模型每次改动都有据可查。希望帮到你。本文还有配套的精品资源点击获取
