Musa Mahmood
8c8b2c6e38
OCR filled the GPU: 9 GB for the mobile model on gpu-bench's 30 images, past 16 GB for the server model, which Windows then paged into shared memory (1.5 to 11 s a capture). onnxruntime's logged arena extensions showed why: the arena keeps every block it takes from CUDA and never joins or frees them, and OCR's inputs grow from run to run (recognizer batches went narrowest first; images in any order across calls), so each larger run left another block behind: 5.7 GB of recognizer blocks where the widest batch needs 0.9 GB. - Detection runs the largest images first, recognition the widest batches. - The detector's and recognizer's arenas shrink once each stage of an ocr_process call is done, by a run on a small input with memory.enable_memory_arena_shrinkage. Shrinking after every run cost 33 to 59% of the speed. - The worker threads are made with the engine and kept, not made for every stage of every call (onnxruntime keeps GPU state for each thread that runs a session, microsoft/onnxruntime#33037); 6% faster. - OCR's arenas grow by what is asked (kSameAsRequested). - ort.log_level sets what the engines' onnxruntime environments log; INFO shows every arena extension. Mobile OCR on gpu-bench (RTX 5080 laptop, 4 workers): peak 5.7 GB over the baseline, 2 GB between calls, 142 ms a capture, against 9.0 GB and 126 ms. The text read differs by a line in 2,986 (batches are grouped differently). The server model is not yet measured again. |
||
|---|---|---|
| .. | ||
| fetch_models.jai | ||
| models.config | ||
| module.jai | ||
| README.md | ||
| tokenizer.jai | ||
Labels
Images and text in one space, from an image-text encoder: SigLIP 2 base at 384 px (Google, Apache
2.0), as converted to ONNX by onnx-community. An image can then be scored against any phrase: a label
list, or words typed into a search. It runs ONNX models through onnxruntime
(modules/onnxruntime), sharing its session and CUDA code with
modules/OCR.
Setting up
cd modules/onnxruntime && jai fetch.jai # if OCR has not fetched it already
cd modules/Labels && jai fetch_models.jai # the vision model (fp32, 373 MB), the text model (int8, 283 MB), tokenizer.json (34 MB)
Each file's SHA-256 is checked. Building the app copies the models into bin/label-models.
Using it
#import "Labels";
engine, ok, error := labels_create(models_folder, "siglip2-base-patch16-384"); // prefer_gpu = true
images := labels_embed_images(engine, pictures); // [] Labels_Image (RGBA): one unit vector each, engine.dimensions long
phrases := labels_embed_texts(engine, .["a photo of a cat", "a screenshot of a terminal"]);
score := labels_score(engine, images.data, phrases.data); // 0 to 1: the model's sigmoid
labels_destroy(engine);
- Images are resized to 384 x 384 (not keeping their shape) with PIL's bilinear filter, which averages over each output pixel's footprint when shrinking, then scaled to -1..1: as SigLIP's own processor does. The vision model runs on CUDA when it can, else on the CPU.
- Text is lowercased (ASCII letters) and tokenized with Gemma's tokenizer (
tokenizer.jai: BPE with byte fallback, read fromtokenizer.json), cut or padded to 64 tokens. The text model runs on the CPU: short inputs, a few at a time. - Scores are
sigmoid(e^logit_scale * cosine + logit_bias), with the trained model'slogit_scaleandlogit_bias(inmodels.config; the ONNX encoders do not have them). They are small for most phrases; their order is what counts. Descriptive prompts ("a screenshot of a code editor showing source code") score far better than single words. labels_to_halfandlabels_from_halfturn a vector into 16-bit floats, as.mcapfiles store them, and back.
Checked against the reference
- The tokenizer gives the same ids as Hugging Face's
tokenizers(0.23.2) on the sametokenizer.json, for 336 strings: label phrases, punctuation, digits, accents, Chinese, Japanese, emoji, tabs and newlines, and text past 64 tokens. - The vectors match SigLIP's own preprocessing through the same ONNX files (Python, onnxruntime and PIL): image vectors at a cosine of 0.9998 or more, text vectors exactly, scores within 0.001.
On the RTX 5080 laptop, an image takes about 18 ms on the GPU (preparing it included) and 190 ms on the CPU; a phrase about 5 ms.