Exporting an Embedding Model to ONNX for Local Inference

Exporting to ONNX is usually pitched as a speed optimisation. The larger practical win is that the resulting artefact runs under onnxruntime with no PyTorch in the image, which changes what you can deploy an embedding model onto.

What ONNX actually buys

ONNX is a serialised computation graph plus weights. Exporting traces the model once and records the operations, which has three consequences worth separating.

• The graph is static. No Python executes per forward pass, so the per-call interpreter overhead disappears. On small models and short texts this overhead is a real fraction of total time, which is why the speedups are largest exactly where the model is smallest.
• The runtime can fuse. onnxruntime can merge attention subgraphs and layer norms into single kernels, which is what the optimisation levels below do.
• The dependency shrinks. A PyTorch wheel with CUDA is measured in gigabytes; onnxruntime for CPU is measured in tens of megabytes. For a container, a Lambda-style function, or a machine that will never have a GPU, that is the deciding factor.

What you lose is dynamism. A traced graph fixes control flow, so anything conditional on input values is baked to whatever branch the trace took. For a standard encoder this is harmless — there is no data- dependent branching in a BERT forward pass — but it is why models with custom modelling code sometimes refuse to export cleanly.

Two shapes must remain dynamic or the export is nearly useless: batch size and sequence length. The exporter marks them as dynamic axes by default, but it is worth confirming, because a graph exported with a fixed sequence length of 512 will pad every input to 512 and quietly throw away the throughput advantage of short documents — which on a corpus of titles and snippets is most of the throughput there was. If your ONNX model is inexplicably slower on short texts than the PyTorch original, this is the first thing to check.

The second thing worth knowing before you start is that the win is concentrated at small scale. The overhead ONNX removes is per-call Python and per-operation dispatch, which is a fixed cost. On a large model with long sequences that cost is a rounding error against the matrix multiplications; on a 33M-parameter model embedding short queries one at a time it can be a third of the wall time. Exporting a large model for batch throughput is often not worth the operational complexity; exporting a small one for a latency-sensitive query path usually is.

Exporting

Sentence Transformers has this built in as a backend, so no separate conversion script is needed:

pip install "sentence-transformers>=3.2" "optimum[onnxruntime]"

python - <<'PY'
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-base-en-v1.5", backend="onnx")
print(model.encode(["hello"]).shape)
model.save_pretrained("bge-base-onnx")
PY

If the repository already contains an ONNX file the backend loads it; if not, it exports on the fly. model_kwargs controls the details: file_name selects which ONNX file to load, defaulting to model.onnx or onnx/model.onnx; provider selects the execution provider, for example CPUExecutionProvider; and export forces the export rather than a load. The Sentence Transformers efficiency documentation lists them.

The saved directory contains the ONNX graph, the tokenizer files and the pooling configuration. That last one matters: the pooling is not inside the ONNX graph, it is a sentence-transformers module applied to the graph’s output. If you plan to run onnxruntime directly without the library, you must reimplement the pooling — CLS or masked mean depending on the family — and getting it wrong is the silent failure described for BGE.

Verifying the output matches

An export that runs is not an export that is correct. The check is numerical and takes ten seconds:

import numpy as np
from sentence_transformers import SentenceTransformer

texts = ["the cat sat on the mat",
         "quarterly revenue rose 4% year over year",
         "SELECT * FROM users WHERE id = 42",
         ""]

torch_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
onnx_model  = SentenceTransformer("BAAI/bge-base-en-v1.5", backend="onnx")

A = torch_model.encode(texts, normalize_embeddings=True)
B = onnx_model.encode(texts, normalize_embeddings=True)

print("max abs diff", np.abs(A - B).max())
print("min cosine  ", (A * B).sum(axis=1).min())
assert np.allclose(A, B, atol=1e-4)

For an fp32 export the cosine should be 1.0 to within floating-point noise and the maximum absolute difference should sit around 1e-5 or better. Anything materially worse means the export changed the computation, and the usual culprits are an opset that lowered an operation differently, a fused kernel with different accumulation order, or — most often — pooling applied differently on the two paths.

Include awkward inputs in the check. An empty string, a text longer than the maximum sequence length, a single emoji, and a batch with very uneven lengths between its members. Padding behaviour is the thing most likely to differ between the two paths, and it only shows up when the batch is ragged.

Optimising and quantizing

Two further steps are available, and they are different operations: graph optimisation preserves the numerics, quantization does not.

from sentence_transformers import SentenceTransformer, export_optimized_onnx_model

model = SentenceTransformer("BAAI/bge-base-en-v1.5", backend="onnx")
export_optimized_onnx_model(model=model, optimization_config="O3",
                            model_name_or_path="bge-base-onnx")

from sentence_transformers import SentenceTransformer, export_dynamic_quantized_onnx_model

model = SentenceTransformer("BAAI/bge-base-en-v1.5", backend="onnx")
export_dynamic_quantized_onnx_model(model=model,
                                    quantization_config="avx512_vnni",
                                    model_name_or_path="bge-base-onnx")

Optimisation levels run O1 through O4, with higher levels fusing more aggressively and O4 including fp16 conversion, which is a numerics change rather than a pure graph rewrite. The quantization configuration names a target instruction set: avx512_vnni requires a CPU with those instructions, and running such a model on a machine without them gives you the size reduction and none of the speed.

The Sentence Transformers documentation publishes its own benchmark results with hardware attached: averaged across four models on an RTX 3090, ONNX gives roughly 1.30x and ONNX-O4 roughly 1.45x; on an Intel i7-13700K, plain ONNX gives roughly 1.20x and int8-quantized ONNX roughly 2.5x. Those are their figures on their hardware, not a measurement made here, and the shape is the transferable part: the quantization win is a CPU story and the graph-optimisation win is a GPU story.

Deploying without PyTorch

The payoff is an image that installs onnxruntime and tokenizers and nothing else. You still need the tokenizer, because ONNX graphs take token ids and not strings, and the tokenizer must be the one the model was trained with.

Running the graph directly means doing pooling and normalisation yourself:

import numpy as np, onnxruntime as ort
from tokenizers import Tokenizer

tok = Tokenizer.from_file("bge-base-onnx/tokenizer.json")
sess = ort.InferenceSession("bge-base-onnx/onnx/model.onnx",
                            providers=["CPUExecutionProvider"])

enc = tok.encode("the cat sat on the mat")
ids = np.array([enc.ids], dtype=np.int64)
mask = np.array([enc.attention_mask], dtype=np.int64)

out = sess.run(None, {"input_ids": ids, "attention_mask": mask,
                      "token_type_ids": np.zeros_like(ids)})[0]

vec = out[:, 0]                              # CLS pooling, for BGE
vec = vec / np.linalg.norm(vec, axis=1, keepdims=True)

The input names are model-specific — print [i.name for i in sess.get_inputs()] rather than assuming three. And re-run the verification from earlier against this hand- written path too, because this is where the pooling reimplementation can quietly diverge.

1. Export with backend="onnx" and save_pretrained.
2. Verify against PyTorch on a batch including empty, over-long and ragged inputs; assert atol=1e-4.
3. Optimise at O2 or O3, or quantize if you are CPU-bound and the target CPU has the instruction set.
4. Re-verify after quantizing, this time expecting drift and measuring top-k overlap rather than cosine.