# How to Export SAM Models to ONNX for Deployment

> Export SAM models to ONNX for deployment using the SamOnnxModel wrapper and export script. Convert prompt encoder and mask decoder efficiently.

- Repository: [Meta Research/segment-anything](https://github.com/facebookresearch/segment-anything)
- Tags: how-to-guide
- Published: 2026-03-07

---

**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`](https://github.com/facebookresearch/segment-anything/blob/main/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__`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/onnx.py#L25-L40) |
| Embed points & masks for tracing | `_embed_points` and `_embed_masks` use only tensor ops, no Python control flow | [`_embed_points`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/onnx.py#L51-L66) & [`_embed_masks`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/onnx.py#L69-L74) |
| Optional single-mask output | `select_masks` rewrites the scoring logic to avoid branching | [`select_masks`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/onnx.py#L92-L104) |
| Stability score replacement | Calls `calculate_stability_score` from [`utils/amg.py`](https://github.com/facebookresearch/segment-anything/blob/main/utils/amg.py) when `use_stability_score=True` | [`forward`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/onnx.py#L27-L34) |

### 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`](https://github.com/facebookresearch/segment-anything/blob/main/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:

```bash
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`](https://github.com/facebookresearch/segment-anything/blob/main/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:

```bash
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:

```python
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`](https://github.com/facebookresearch/segment-anything/blob/main/scripts/export_onnx_model.py).

## Key Source Files

| File | Role | Link |
|------|------|------|
| [`segment_anything/utils/onnx.py`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/onnx.py) | Defines `SamOnnxModel`, the ONNX‑exportable wrapper combining prompt encoder and mask decoder. | [SamOnnxModel source](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/onnx.py) |
| [`scripts/export_onnx_model.py`](https://github.com/facebookresearch/segment-anything/blob/main/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](https://github.com/facebookresearch/segment-anything/blob/main/scripts/export_onnx_model.py) |
| [`segment_anything/__init__.py`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/__init__.py) | Provides `sam_model_registry` mapping model‑type strings to constructors. | [Registry source](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/__init__.py) |
| [`segment_anything/utils/amg.py`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/amg.py) | Implements `calculate_stability_score` used when `--use-stability-score` is set. | [AMG utilities](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/amg.py) |
| [`segment_anything/modeling/sam.py`](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/sam.py) | Core SAM model definition (image encoder, prompt encoder, mask decoder). | [SAM model source](https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/sam.py) |

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`](https://github.com/facebookresearch/segment-anything/blob/main/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`](https://github.com/facebookresearch/segment-anything/blob/main/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`](https://github.com/facebookresearch/segment-anything/blob/main/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`](https://github.com/facebookresearch/segment-anything/blob/main/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`](https://github.com/facebookresearch/segment-anything/blob/main/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`](https://github.com/facebookresearch/segment-anything/blob/main/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.