---
license: apache-2.0
base_model: google/siglip2-so400m-patch16-naflex
tags:
  - onnx
  - tensorrt
  - siglip2
  - naflex
  - vision
  - inference-optimization
library_name: onnx
pipeline_tag: feature-extraction
---

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

An ONNX export of the **vision-encoder tail** of
[`google/siglip2-so400m-patch16-naflex`](https://huggingface.co/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

```python
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

```bash
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`](https://huggingface.co/google/siglip2-so400m-patch16-naflex).
All weights are Google's; this repository contributes only a graph
transformation. Please cite the original work:

```bibtex
@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](https://huggingface.co/guruansh).