guruansh/siglip2-so400m-naflex-onnx

SigLIP2 so400m NaFlex — vision encoder tail, ONNX (fp16)

An ONNX export of the vision-encoder tail of google/siglip2-so400m-patch16-naflex, built so the model can be compiled with TensorRT without breaking NaFlex's native-aspect-ratio handling.

Benchmarked on four GPUs. Up to 2.08× faster than PyTorch fp16, with no measurable accuracy loss.

---

⚠️ Read this first — this is NOT a standalone model

This graph contains the 27 transformer layers + post-layernorm + attention-pooling head only. It does not contain:

• the NaFlex embeddings front-end (patch embedding + position-embedding interpolation) — deliberately left in PyTorch, see below
• the text tower

Its input is not an image. It is the tensor produced by vision_model.embeddings(pixel_values, spatial_shapes). Feeding images directly will fail, and guessing the interface silently produces wrong numbers.

Why the model is split here

NaFlex preserves each image's native aspect ratio: a tall narrow crop gets a patch grid like [36, 16], a square one gets [24, 24]. Crucially, the tensor shape is identical either way — pixel_values is always (B, 576, 768). What varies is the value of spatial_shapes, which drives a per-image interpolation of the learned position-embedding grid.

That interpolation is data-dependent control flow — a Python loop calling F.interpolate with a size read out of the tensor. Tracing it bakes in whichever aspect ratios happened to be in the export batch, so every other aspect ratio then silently receives the wrong position embeddings. The model still runs, still returns finite, normalised, plausible-looking vectors. Nothing errors.

Keeping the embeddings in PyTorch and exporting only the static tail avoids this entirely. It costs almost nothing: the embeddings block measured 2.7% of total runtime.

Verification. Two gates, both in export_tail.py (--verify):

check	result
split path vs full get_image_features	max abs diff 0.00000000, cosine 0.99999988
ONNX vs PyTorch on aspect ratios not in the export batch	cosine min 0.999577

The second is the one that matters — a baked-in interpolation passes on traced shapes and fails on held-out ones.

I/O contract

	name	dtype	shape
input	hidden_states	float16	[batch, 576, 1152]
input	attention_mask	int32	[batch, 576]
output	pooled	float16	[batch, 1152]

opset 17 · 854 MB · batch is dynamic · all 445 initializers are fp16

attention_mask is the processor's pixel_attention_mask (1 = real patch, 0 = padding; 551–576 of the 576 slots are typically real).

Bind these dtypes exactly. TensorRT's set_tensor_address takes a raw pointer and performs no dtype checking — binding fp32 buffers to this engine yields normal-looking latencies and non-finite outputs, with no error raised. Query engine.get_tensor_dtype() rather than assuming.

Usage

import numpy as np, torch, onnxruntime as ort
from transformers import AutoModel, AutoProcessor
from PIL import Image

MODEL_ID, REV = "google/siglip2-so400m-patch16-naflex", "cc24074f717b612951c2dead130904ab9b65a81e"
MAX_NUM_PATCHES = 576                      # must match the export

proc  = AutoProcessor.from_pretrained(MODEL_ID, revision=REV)
torch_model = AutoModel.from_pretrained(MODEL_ID, revision=REV,
                                        torch_dtype=torch.float16).cuda().eval()

images = [Image.open("a.jpg").convert("RGB")]
inp = proc(images=images, return_tensors="pt", max_num_patches=MAX_NUM_PATCHES)
inp = {k: v.cuda() for k, v in inp.items()}

# --- front end stays in PyTorch: this is where NaFlex lives ---
with torch.no_grad(), torch.autocast("cuda"):
    hidden = torch_model.vision_model.embeddings(inp["pixel_values"], inp["spatial_shapes"])

# --- tail runs from this ONNX ---
sess = ort.InferenceSession("siglip2_tail_fp16.onnx", providers=["CUDAExecutionProvider"])
pooled = sess.run(["pooled"], {
    "hidden_states":  hidden.half().cpu().numpy(),
    "attention_mask": inp["pixel_attention_mask"].to(torch.int32).cpu().numpy(),
})[0]

emb = pooled / np.linalg.norm(pooled, axis=-1, keepdims=True)   # L2 normalise

Building a TensorRT engine

trtexec --onnx=siglip2_tail_fp16.onnx \
        --saveEngine=engine.plan \
        --shapes=hidden_states:8x576x1152,attention_mask:8x576

Builds in 51–61 s across the GPUs tested.

A TensorRT .plan is compiled for one GPU architecture, one TensorRT version, and one batch profile. No engine is shipped here on purpose — this ONNX is portable, and each machine should build its own. On TensorRT ≥ 11 the network is strongly typed: precision comes from the ONNX itself, so this fp16 graph yields an fp16 engine without any builder flag.

Measured performance

SigLIP2 vision tower alone, CUDA-event timed, inputs resident on GPU, max_num_patches=576. "crops/s" = images encoded per second.

GPU	arch	PyTorch fp16	+ torch.compile	TensorRT (end-to-end)	TensorRT (engine only)	speedup
Tesla T4	sm_75	33.7	36.9	36.8	37.3	1.09×
NVIDIA L4	sm_89	67.0	70.3	89.7	91.1	1.34×
RTX 5070 Ti	sm_120	52.9	57.2	109.8	115.0	2.08×
RTX 5090	sm_120	232.9	263.4	457.1	497.2	1.96×

"end-to-end" includes the PyTorch NaFlex front end; "engine only" is the ONNX graph alone.

The benefit is strongly architecture-dependent. On T4 TensorRT merely ties torch.compile; on consumer Blackwell it roughly doubles throughput. Benchmark on your own target rather than assuming. Note also that the RTX 5070 Ti is slower than an L4 in plain PyTorch fp16 (52.9 vs 67.0) yet faster once compiled with TensorRT (109.8 vs 89.7).

Stacks: torch 2.7.0–2.7.1 + cu128 · TensorRT 11.2.1.2 · transformers 5.15.1

Accuracy

Evaluated on 1016 person crops (COCO-derived) against 21,450 image–prompt similarity scores, versus the PyTorch fp16 reference (L4 engine):

	value
embedding cosine vs PyTorch fp16	mean 0.999736, min 0.997570
mean absolute change in similarity score	0.00058 (scores span ~0–0.24)
signed mean change	+0.000003 — no systematic bias
ranking accuracy (best positive vs best negative prompt)	unchanged, ±0.00 pp

Per-batch sanity checks on the other GPUs agreed: cosine min 0.999967 (RTX 5090) and 0.999974 (RTX 5070 Ti) against PyTorch on identical inputs.

This is fp16 numerics throughout — not quantisation — so the residual difference is only kernel selection and accumulation order.

For calibration: running the same fp16 maths on two different GPUs already moves the scores by 0.000088 on average. The TensorRT difference is small but above that floor.

Limitations

• Tail only. Requires the PyTorch embeddings front-end (see Usage).
• Fixed at max_num_patches=576. A different patch budget needs a re-export.
• Exported from revision cc24074f... of the base model.
• No text tower — encode prompts with the original model and cache them.
• Not evaluated beyond person crops; the base model's own limitations (small objects, attribute binding, counting, gaze, expression) are unchanged by this export.

Changelog

• v2 — re-exported natively in fp16. The first upload was produced by converting an fp32 export with onnxconverter_common.float16, which left some constants as Float; TensorRT ≥ 11 is strict about mixed types and refused to parse it (ElementWiseOperation SUB must have same input types). This version traces directly in fp16, so all 445 initializers are fp16 and the graph builds cleanly. Added RTX 5070 Ti and RTX 5090 benchmarks.
• v1 — initial release (T4 and L4 benchmarks). Do not use — TensorRT cannot parse it.

License and attribution

Apache-2.0, inherited from google/siglip2-so400m-patch16-naflex. All weights are Google's; this repository contributes only a graph transformation. Please cite the original work:

@article{tschannen2025siglip2,
  title   = {SigLIP 2: Multilingual Vision-Language Encoders with Improved
             Semantic Understanding, Localization, and Dense Features},
  author  = {Tschannen, Michael and others},
  journal = {arXiv preprint arXiv:2502.14786},
  year    = {2025}
}

Export and benchmarking by @guruansh.