How to Export SAM Models to ONNX for Deployment

Export SAM models to ONNX by using the SamOnnxModel wrapper class and the provided export script to convert the prompt encoder and mask decoder, while caching the image encoder output separately.

The Segment Anything Model (SAM) from Meta's facebookresearch/segment-anything repository provides powerful zero-shot segmentation capabilities. To deploy SAM in production environments with optimized inference, you need to export SAM models to ONNX format. This process involves isolating the prompt encoder and mask decoder components while handling the image encoder separately.

Understanding SAM Architecture for ONNX Export

SAM consists of three core components that typically run together during inference. When you export SAM models to ONNX, understanding these components is crucial for optimizing the deployment pipeline.

Core Components

1. Image Encoder – extracts a dense visual embedding from the input image. This is typically a Vision Transformer (ViT) backbone.
2. Prompt Encoder – embeds user prompts such as points, boxes, or masks into the model's embedding space.
3. Mask Decoder – predicts segmentation masks from the image embedding and prompt embeddings.

Why Export Only the Prompt Encoder and Mask Decoder

When exporting to ONNX, you only need the prompt encoder and mask decoder because the image encoder is typically run once per image and its output (image_embeddings) can be cached. This separation significantly reduces the ONNX model size and allows you to precompute image embeddings for static images, making interactive prompt-based inference much faster.

The SamOnnxModel Wrapper

The repository provides a thin wrapper class SamOnnxModel in segment_anything/utils/onnx.py that combines the prompt encoder and mask decoder and adds the necessary post-processing steps. The wrapper is deliberately written to be trace-friendly for ONNX conversion.

Key Implementation Details

The SamOnnxModel class implements several trace-friendly modifications:

Feature	Implementation	Source
Combine prompt encoder & mask decoder	self.mask_decoder = model.mask_decoder and calls to self.model.prompt_encoder	SamOnnxModel.__init__
Embed points & masks for tracing	_embed_points and _embed_masks use only tensor ops, no Python control flow	_embed_points & _embed_masks
Optional single-mask output	select_masks rewrites the scoring logic to avoid branching	select_masks
Stability score replacement	Calls calculate_stability_score from utils/amg.py when use_stability_score=True	forward

Output Options and Post-Processing

The wrapper supports several optional outputs controlled by constructor flags. The resulting ONNX file contains three (or up to five) outputs depending on the flags:

Output name	Meaning
masks	Upscaled binary masks (H × W)
iou_predictions	Predicted quality scores (or stability scores if --use-stability-score)
low_res_masks	Low‑resolution logits returned for further processing
stability_scores (optional)	Stability‑score per mask
areas (optional)	Pixel area of each mask
low_res_logits (optional)	Raw low‑resolution mask logits

Exporting SAM to ONNX

The scripts/export_onnx_model.py script provides a complete CLI utility for exporting SAM models to ONNX format.

Using the Export Script (CLI)

Run the export script with your checkpoint and desired options:

python -m segment_anything.scripts.export_onnx_model \
    --checkpoint /path/to/sam_vit_h_4b8939.pth \
    --model-type vit_h \
    --output sam_vit_h.onnx \
    --return-single-mask \
    --use-stability-score \
    --gelu-approximate \
    --opset 17

Export Parameters Explained

• --checkpoint: Path to a pretrained SAM .pth file.
• --model-type: Architecture variant (default, vit_h, vit_l, or vit_b).
• --return-single-mask: Makes the ONNX model output only the best mask, reducing post-processing cost.
• --use-stability-score: Swaps the default IoU prediction for a stability‑score‑based metric from segment_anything/utils/amg.py.
• --gelu-approximate: Replaces the costly erf‑based GELU with a tanh approximation, improving compatibility with runtimes that lack erf support.
• --opset: ONNX opset version (default 17).

Quantization for CPU Deployment

For optimized CPU inference, you can quantize the exported model using the same script:

python -m segment_anything.scripts.export_onnx_model \
    --checkpoint /path/to/sam_vit_h_4b8939.pth \
    --model-type vit_h \
    --output sam_vit_h.onnx \
    --quantize-out sam_vit_h_quantized.onnx \
    --opset 17

The script invokes onnxruntime.quantization.quantize_dynamic to produce a smaller, integer‑only model that runs faster on CPU.

Running Inference with the ONNX Model

After exporting, you can run inference using ONNX Runtime. The input shapes must match the dummy inputs used during export:

import numpy as np
import onnxruntime as ort
import torch

# Prepare inputs matching the export configuration

image_embeddings = torch.randn(1, 256, 64, 64)        # Adjust dims for your model

point_coords = torch.tensor([[[300.0, 400.0]]])      # (batch, num_points, 2)

point_labels = torch.tensor([[1]])                   # (batch, num_points)

mask_input = torch.randn(1, 1, 256, 256)             # Low-res mask prompt

has_mask_input = torch.tensor([1.0])
orig_im_size = torch.tensor([1024.0, 768.0])

# Convert to NumPy for ONNX Runtime

np_inputs = {
    "image_embeddings": image_embeddings.numpy(),
    "point_coords": point_coords.numpy(),
    "point_labels": point_labels.numpy(),
    "mask_input": mask_input.numpy(),
    "has_mask_input": has_mask_input.numpy(),
    "orig_im_size": orig_im_size.numpy(),
}

# Run inference

session = ort.InferenceSession("sam_vit_h.onnx", providers=["CPUExecutionProvider"])
outputs = session.run(None, np_inputs)

masks, scores, low_res = outputs[:3]   # Adjust if you used extra-metrics flags

print("Masks shape:", masks.shape)
print("Scores shape:", scores.shape)

This snippet mirrors the verification step performed by the export script in scripts/export_onnx_model.py.

Key Source Files

File	Role	Link
segment_anything/utils/onnx.py	Defines SamOnnxModel, the ONNX‑exportable wrapper combining prompt encoder and mask decoder.	SamOnnxModel source
scripts/export_onnx_model.py	CLI utility that loads a checkpoint, builds the wrapper, and runs torch.onnx.export. Handles optional quantization.	Export script source
segment_anything/__init__.py	Provides sam_model_registry mapping model‑type strings to constructors.	Registry source
segment_anything/utils/amg.py	Implements calculate_stability_score used when --use-stability-score is set.	AMG utilities
segment_anything/modeling/sam.py	Core SAM model definition (image encoder, prompt encoder, mask decoder).	SAM model source

These files together form the complete pipeline for turning a pretrained SAM checkpoint into an ONNX model ready for deployment.

Summary

• Export SAM models to ONNX by isolating the prompt encoder and mask decoder using the SamOnnxModel wrapper in segment_anything/utils/onnx.py.
• The image encoder is excluded from ONNX export because its output can be cached per image, significantly reducing model size and enabling faster interactive inference.
• Use the provided scripts/export_onnx_model.py CLI tool with flags like --return-single-mask, --use-stability-score, and --gelu-approximate to customize the export for your deployment target.
• Enable dynamic axes for point_coords and point_labels to support variable numbers of prompt points during inference.
• Optionally quantize the exported ONNX model using --quantize-out for optimized CPU performance with ONNX Runtime.

Frequently Asked Questions

Why is the image encoder not included in the ONNX export?

The image encoder is excluded because it is computationally expensive and only needs to run once per image. By caching the image_embeddings output and exporting only the prompt encoder and mask decoder to ONNX, you create a lightweight model that enables real-time interactive segmentation. This architecture allows users to change prompts rapidly without reprocessing the image through the heavy ViT backbone.

What is the difference between IoU predictions and stability scores?

IoU predictions are the default quality estimates output by the mask decoder, representing the model's confidence in each predicted mask. Stability scores, enabled with the --use-stability-score flag, are computed by calculate_stability_score in segment_anything/utils/amg.py and measure mask quality by comparing predictions at different threshold levels. Stability scores often provide more reliable rankings for ambiguous prompts and complex boundaries.

Can I export SAM to ONNX without using the provided script?

Yes, you can export SAM to ONNX programmatically by instantiating SamOnnxModel from segment_anything/utils/onnx.py and calling torch.onnx.export directly. You will need to construct dummy inputs matching the expected shapes for image_embeddings, point_coords, point_labels, mask_input, has_mask_input, and orig_im_size, and specify dynamic axes for the point dimensions. However, using scripts/export_onnx_model.py is recommended as it handles GELU approximation, stability score integration, and automatic verification.

How do I handle dynamic input sizes for point prompts?

The export script configures dynamic axes for the point_coords and point_labels inputs, allowing the ONNX model to accept varying numbers of prompt points. In scripts/export_onnx_model.py, dynamic axes are defined as {1: "num_points"} for both tensors, enabling the batch dimension to remain fixed while the point count varies. When running inference, simply provide NumPy arrays with the desired number of points in the second dimension, and ONNX Runtime will handle the variable shape automatically.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}