Segment Anything 2 on Jetson Thor

Take Meta's SAM 2.1 from a Hugging Face checkpoint to accelerated engines for NVIDIA Jetson Thor, with each precision priced on your own images. Along the way you will meet everything the Physical AI lane cares about — promptable models that split in two, sample sets that double as calibration data, and a scorecard that answers what quantization costs you rather than what it costs in general.

Prerequisites

• A Melange account enrolled in the Early Access Program, and a personal key (how to get one)
• Python 3.10+ and pip install torch torchvision transformers onnx onnxscript onnxruntime numpy pillow — torchvision is what Sam2Processor loads its image processor from, and onnxscript is what the dynamo exporter runs on. Neither is pulled in automatically, and both fail at import time rather than at export time. This walkthrough was run on torch 2.13 and transformers 5.15.
• A folder of representative images — around eight is enough
• An NVIDIA Jetson Thor device, once you are ready to deploy — the export, upload, and scorecard steps need no device

How SAM 2 splits

SAM 2 is promptable: the expensive part looks at the image, the cheap part answers prompts about it. That split is the reason it becomes two models, not one.

Uploading them separately is what lets you encode an image once and then answer clicks at interactive rates. It also means two model keys, two sample sets, and two scorecards.

Your application keeps the pre- and post-processing: resize to 1024×1024 and normalize on the way in, and upsample the decoder's 256×256 low-resolution mask logits back to the original image size on the way out.

Step 1: Export both modules to ONNX

Wrap the encoder and decoder as nn.Modules and export each with the dynamo exporter.

import torch
from PIL import Image
from transformers import Sam2Model, Sam2Processor

REPO = "facebook/sam2.1-hiera-small"
OPSET = 18

model = Sam2Model.from_pretrained(REPO).eval()
processor = Sam2Processor.from_pretrained(REPO)

image = Image.open("example.jpg").convert("RGB")
inputs = processor(images=image, input_points=[[[[400, 650]]]], return_tensors="pt")
pixel_values = inputs["pixel_values"].float()
point_coords = inputs["input_points"].float()
# 1 marks a foreground click; the processor does not emit labels on its own.
point_labels = torch.ones(point_coords.shape[:-1], dtype=torch.int32)


class Encoder(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, pixel_values):
        feat_s0, feat_s1, image_embed = self.model.get_image_embeddings(pixel_values)
        return feat_s0, feat_s1, image_embed


class Decoder(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, image_embed, feat_s0, feat_s1, point_coords, point_labels):
        positional_embeddings = self.model.get_image_wide_positional_embeddings()
        sparse, dense = self.model.prompt_encoder(
            input_points=point_coords,
            input_labels=point_labels,
            input_boxes=None,
            input_masks=None,
        )
        masks, iou_scores, _, _ = self.model.mask_decoder(
            image_embeddings=image_embed,
            image_positional_embeddings=positional_embeddings,
            sparse_prompt_embeddings=sparse,
            dense_prompt_embeddings=dense,
            multimask_output=True,
            high_resolution_features=[feat_s0, feat_s1],
        )
        return masks, iou_scores


torch.onnx.export(
    Encoder(model).eval(),
    (pixel_values,),
    "sam2_encoder.onnx",
    input_names=["pixel_values"],
    output_names=["feat_s0", "feat_s1", "image_embed"],
    opset_version=OPSET,
    dynamo=True,
)

with torch.no_grad():
    feat_s0, feat_s1, image_embed = Encoder(model)(pixel_values)

num_points = torch.export.Dim("num_points", min=1, max=64)
torch.onnx.export(
    Decoder(model).eval(),
    (image_embed, feat_s0, feat_s1, point_coords, point_labels),
    "sam2_decoder.onnx",
    input_names=["image_embed", "feat_s0", "feat_s1", "point_coords", "point_labels"],
    output_names=["masks", "iou_scores"],
    dynamic_shapes=(None, None, None, {2: num_points}, {2: num_points}),
    opset_version=OPSET,
    dynamo=True,
)

Each export writes two files: the graph, and a .onnx.data sidecar holding the weights — sam2_encoder.onnx (0.7 MB) next to sam2_encoder.onnx.data (137 MB), and the same pair for the decoder. Both are needed; the graph alone is not a model. Step 3 uploads them together.

Step 2: Build the sample sets

Each module needs its own set. The encoder's samples are preprocessed images. The decoder's samples are the encoder's real outputs for those same images, plus the prompt — feeding it anything else would calibrate it on a distribution it never sees.

from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image
from transformers import Sam2Processor

processor = Sam2Processor.from_pretrained("facebook/sam2.1-hiera-small")
session = ort.InferenceSession("sam2_encoder.onnx", providers=["CPUExecutionProvider"])

# One image and one click per sample, in original pixel coordinates.
picks = [
    ("images/000000039769.jpg", (400, 650)),
    ("images/000000000285.jpg", (300, 400)),
    # ... eight or so, covering the variety your application will see
]

Path("samples/encoder").mkdir(parents=True, exist_ok=True)
Path("samples/decoder").mkdir(parents=True, exist_ok=True)

for index, (path, (x, y)) in enumerate(picks):
    image = Image.open(path).convert("RGB")
    inputs = processor(images=image, input_points=[[[[x, y]]]], return_tensors="pt")
    pixel_values = inputs["pixel_values"].float().numpy()

    np.savez_compressed(
        f"samples/encoder/sample_{index:02d}.npz",
        pixel_values=pixel_values,
    )

    feat_s0, feat_s1, image_embed = session.run(
        ["feat_s0", "feat_s1", "image_embed"],
        {"pixel_values": pixel_values},
    )
    point_coords = inputs["input_points"].float().numpy()
    point_labels = np.ones(point_coords.shape[:-1], dtype=np.int32)

    np.savez_compressed(
        f"samples/decoder/sample_{index:02d}.npz",
        image_embed=image_embed,
        feat_s0=feat_s0,
        feat_s1=feat_s1,
        point_coords=point_coords,
        point_labels=point_labels,
    )

Two details that are easy to get wrong and are both rejected at upload:

• The processor rescales your click into the model's 1024×1024 space. Pass original-image coordinates and let it do that — hand-scaling is where the pixel-offset bugs live.
• point_labels must be int32, matching the export above. NumPy would give you int64 by default.

Run the validation script against each pair before uploading — the model and its own sample directory, never the other module's:

python validate_samples.py sam2_encoder.onnx samples/encoder
python validate_samples.py sam2_decoder.onnx samples/decoder

Step 3: Upload as two models

Repeat this once per module:

1. Log in to the Melange Dashboard and create a project.
2. Click Upload and choose Physical AI Model as the model type.
3. Attach sam2_encoder.onnx under Model File and sam2_encoder.onnx.data under External Data Files. The dialog reads the graph and names the sidecar it expects.
4. Attach the samples/encoder/*.npz archives under Calibration Sample Files.
5. Start the run and wait for the status to reach Ready.

Then do the same with sam2_decoder.onnx, its sam2_decoder.onnx.data, and samples/decoder/. You end up with two model keys, for example your-account/sam2-encoder and your-account/sam2-decoder.

Step 4: Read the Benchmark scorecard

When the run finishes, the model report carries a Benchmark table with one row per precision:

Output SNR is the column to read first. It answers one question: how far did this engine drift from the model you uploaded? Higher is closer, and ∞ means bit-exact. Melange bands it for you — High fidelity at 40 dB and above, Moderate from 20 dB, Low below that.

Two things worth internalizing before you pick a row:

• The score is fidelity, not accuracy. It measures agreement with your own FP32 model, not correctness on some ground truth. An image your original model segments badly will be segmented badly by every engine in the table, at any SNR.
• Cheaper is not automatically faster. Quantization adds conversion work at precision boundaries, and on some graphs that costs more than the arithmetic it saves. The latency column is measured, not predicted — read it rather than assuming the ladder goes down.

The encoder and decoder are scored independently, which is the point of splitting them: they tolerate quantization differently, and you can ship one at a lower precision than the other.

Step 5: Run on Jetson Thor

Everything up to here is self-serve. The scorecard from Step 4 already answers the question you would otherwise answer by trial on the device: which precision to ship, measured on your own images.

How the two modules chain

Whatever you load them with, the split fixes the shape of the application:

• Run the encoder once per frame and hold its three outputs — feat_s0, feat_s1, and image_embed.
• Run the decoder once per click, passing those three tensors plus point_coords and point_labels, in the module's declared input order.
• The decoder returns three candidate masks as 256×256 logits, each with an IoU score. Take the highest-scoring one, threshold it at zero, and upsample it to your original image size.

That is the whole point of uploading two models: the expensive half runs once per image, and clicks stay interactive.

Related resources