在 Devin 中试用 DeepWiki 处理私有仓库
DeepWiki

facebookresearch/segment-anything

菜单

ONNX 导出

相关源文件

本页面介绍了将 SAM (Segment Anything Model) 模型导出为 ONNX 格式以及使用导出模型的流程。ONNX 导出功能允许将 SAM 的提示编码器和掩码解码器组件部署到支持 ONNX Runtime 的不同平台，同时排除计算量大的图像编码器，后者最好单独运行。

有关核心 SAM 架构的信息，请参阅 SAM 模型架构。有关不使用 ONNX 的模型使用详情，请参阅 SamPredictor。

ONNX 导出概述

ONNX (Open Neural Network Exchange) 是表示深度学习模型的一种开放格式。SAM 的 ONNX 导出功能能够：

  1. 在包括 Web 浏览器在内的各种平台上运行提示编码器和掩码解码器。
  2. 将图像嵌入计算（计算量较大）与交互式提示分离。
  3. 优化掩码生成，以便在资源受限的环境中进行推理。

来源

导出到 ONNX 的组件

ONNX 导出过程仅导出 SAM 模型中的特定组件。这种架构决策能够实现高效部署，同时保持性能。

导出到 ONNX 的组件包括：

  • 提示编码器：处理输入提示（点、框或现有掩码）。
  • 掩码解码器：根据图像嵌入和编码的提示生成掩码预测。
  • 掩码后处理：调整掩码大小并准备最终掩码。

计算量大的图像编码器 (ViT) 已从 ONNX 导出中排除，因为它通常每个图像只运行一次，并且单独在 GPU 上执行效率更高。

来源

导出为 ONNX 格式

可以使用导出脚本或以编程方式将 SAM 模型导出为 ONNX 格式。导出过程将创建一个 ONNX 模型，该模型以图像嵌入和提示作为输入，并输出预测的掩码。

使用导出脚本

将 SAM 模型导出为 ONNX 的主要方法是使用导出脚本：

python scripts/export_onnx_model.py \
  --checkpoint path/to/sam_checkpoint.pth \
  --output sam_onnx_model.onnx \
  --model-type vit_h \
  --return-single-mask
选项描述默认
--checkpointSAM 模型检查点的路径。必填
--output保存 ONNX 模型的路径。必填
--model-typeSAM 模型类型（vit_h、vit_l、vit_b）。必填
--return-single-mask仅返回最佳掩码，而不是多个掩码。否
--opset要使用的 ONNX opset 版本（必须 ≥11）。17
--quantize-out保存量化模型的路径（如果指定）。None
--gelu-approximate用 tanh 近似替换 GELU。否
--use-stability-score使用稳定性分数而不是预测质量。否
--return-extra-metrics返回额外的指标（高分辨率时较慢）。否

来源

通过编程方式导出

为了获得更多控制权，您可以通过编程方式导出模型：

来源

量化

您可以选择对导出的 ONNX 模型进行量化，以减小其大小并提高推理速度。这对于 Web 部署特别有用。

来源

使用 ONNX 模型

使用导出的 ONNX 模型需要：

  1. 使用 SAM 图像编码器计算图像嵌入。
  2. 以正确格式准备提示。
  3. 使用适当的输入运行 ONNX 模型。
  4. 后处理输出掩码。

输入格式

ONNX 模型需要的输入格式与标准的 SamPredictor 不同。

输入描述形状类型
image_embeddings图像编码器的输出。[1, embed_dim, embed_h, embed_w]float32
point_coords在 1024×1024 空间中的点/框坐标。[1, num_points, 2]float32
point_labels点的标签（0=负，1=正，2=左上角框，3=右下角框，-1=填充）。[1, num_points]float32
mask_input前一个掩码作为输入（如果没有则为零）。[1, 1, 256, 256]float32
has_mask_input指示是否使用掩码输入（0/1）。[1]float32
orig_im_size原始图像尺寸（高，宽）。[2]float32

模型产生三个输出：

  • masks：预测的掩码，形状为 [1, 1, H, W]。
  • iou_predictions：预测的 IoU 分数。
  • low_res_masks：低分辨率掩码 logits。

来源

工作流程示例

使用 ONNX 模型的一般工作流程包括：

来源

代码示例

下面是一个使用 ONNX 模型的简化示例：

来源

性能考量

在使用 ONNX 导出的模型时，请考虑以下性能方面：

  1. 图像编码器分离：图像编码器（最重部分）已从 ONNX 导出中排除，应单独运行，通常在 GPU 上。

  2. 量化优势：量化模型可以显著减小尺寸并提高推理速度，尤其是在 Web 部署中。

  3. 单个或多个掩码：return_single_mask 选项通过避免对多个掩码候选进行放大来提高高分辨率图像的性能。

  4. GELU 近似：gelu-approximate 选项可以提高在具有缓慢或未实现 erf 操作的运行时上的性能。

  5. 额外指标：返回额外指标会减慢高分辨率输出的处理速度。

ONNX 导出过程专门设计用于实现高效的交互式分割，其中图像只需处理一次，然后各种提示可以快速处理，而无需重新计算图像嵌入。

来源

部署场景

ONNX 导出的模型在以下场景中特别有用：

  1. Web 应用：轻量级的 ONNX 模型可以使用 ONNX Web Runtime 在浏览器中运行，而图像编码器可以在服务器上或使用 WebGL 运行。

  2. 移动设备：组件的分离允许在资源受限的设备上高效部署。

  3. 多框架环境：使用多个 ML 框架的组织可以将 SAM 集成到其现有流程中。

  4. 边缘设备：量化模型可以在计算资源有限的边缘设备上高效运行。

来源

技术实现细节

ONNX 导出功能围绕 SamOnnxModel 类构建，该类封装了标准的 SAM 模型并进行必要的适应以实现 ONNX 兼容性。

SamOnnxModel 结构

SamOnnxModel 类包含：

  1. 点嵌入：调整提示编码器的点嵌入以进行 ONNX 追踪。
  2. 掩码嵌入：处理模型的掩码输入。
  3. 掩码后处理：负责将掩码调整到原始图像尺寸。
  4. 掩码选择：当启用了 return_single_mask 时，实现选择最佳掩码的逻辑。

主要实现细节

该实现避免了可能与 ONNX 导出不兼容的控制流操作，并使用了可以正确追踪的张量操作。

来源

常见问题排查

在导出或使用 ONNX 格式的 SAM 时，您可能会遇到以下常见问题：

  1. ONNX 版本兼容性：确保 ONNX opset 版本（默认为 17）与您的运行时环境兼容。如有必要，请使用 --opset 参数选择不同的版本。

  2. 输入准备：ONNX 模型需要特定格式的输入。大多数错误都与不正确的输入形状或类型有关。

  3. 点坐标：请记住，在将点坐标传递给 ONNX 模型之前，必须将其转换为模型 1024×1024 的坐标空间。

  4. 点填充：对于框输入或在未提供点的情况下，必须包含标签为 -1 的填充点。

  5. 量化问题：如果量化模型产生不正确的结果，请尝试使用非量化版本，它可能更准确但速度较慢。

来源

结论

SAM 中的 ONNX 导出功能提供了一种灵活的方式，可以将模型交互式分割功能部署到各种平台。通过将一次性的图像嵌入计算与交互式提示过程分离开来，它能够为实时分割应用实现高效的工作流程。