--- license: apache-2.0 library_name: litert pipeline_tag: image-feature-extraction base_model: timm/vit_pe_core_base_patch16_224.fb tags: - litert - tflite - on-device - android - gpu - clip - perception-encoder - image-encoder - vit - rope --- # Perception Encoder (PE-Core-B16-224) — LiteRT (TFLite) GPU On-device [LiteRT](https://ai.google.dev/edge/litert) (`.tflite`) conversion of **Perception Encoder Core** (PE-Core, Meta 2025), the SOTA CLIP-style image tower, converted from [`timm/vit_pe_core_base_patch16_224.fb`](https://huggingface.co/timm/vit_pe_core_base_patch16_224.fb) (ViT-B/16, 94M params; original [facebook/PE-Core-B16-224](https://huggingface.co/facebook/PE-Core-B16-224)). A single forward pass turns one RGB image into a **1024-d L2-normalized image embedding** for zero-shot classification, retrieval, and similarity — running **fully on the LiteRT `CompiledModel` GPU accelerator** (ML Drift): **all 1028 ops are GPU-native (`Replacing 1028 out of 1028 node(s) ... LITERT_CL`), no CPU fallback, no Flex ops.** ## Files | File | Size | Description | |------|------|-------------| | `pe_core_base_224_fp16.tflite` | 187 MB | FP16 single-graph model, GPU full-residency | | `convert_pecore.py` | — | Reproducible conversion script (timm → tflite) | ## I/O - **Input**: `[1, 3, 224, 224]` float32, **NCHW**, RGB normalized to **`[-1, 1]`** i.e. `(pixel/255 - 0.5) / 0.5` (timm mean/std = `(0.5, 0.5, 0.5)`). Normalization is applied by the caller (not baked into the graph). - **Output**: `[1, 1024]` float32, **L2-normalized** image embedding. ## Usage (Android, LiteRT CompiledModel) ```kotlin val model = CompiledModel.create( context.assets, "pe_core_base_224_fp16.tflite", CompiledModel.Options(Accelerator.GPU), null ) val inputs = model.createInputBuffers() val outputs = model.createOutputBuffers() inputs[0].writeFloat(nchwFloatArray) // [1,3,224,224], RGB scaled to [-1,1] model.run(inputs, outputs) val embedding = outputs[0].readFloat() // [1024], already L2-normalized ``` For zero-shot classification, precompute text-label embeddings with the PE-Core text tower offline and take the dot product on device. ## Performance - **~66 ms / image steady-state** on a Pixel 8a (Mali-G615) GPU (best 12.5 ms), full GPU residency, FP16. ## Conversion notes Converted with [litert-torch / ai-edge-torch](https://github.com/google-ai-edge/ai-edge-torch). Making a RoPE ViT image tower **fully GPU-resident *and* numerically correct** on the ML Drift GPU delegate required four verbatim (weights-exact, output corr ≈ 1.0) model-side rewrites — the first three for residency, the last for on-device numerical correctness: 1. **Fused-qkv → 4D manual attention** — the fused `qkv` reshape emits a 5D head-split the GPU delegate rejects; decompose into separate q/k/v projections. Self-attention uses `scaled_dot_product_attention`, whose lowering keeps the batch-matmul 3D with a materialized transpose (both required for residency). 2. **Interleaved 2D-RoPE → rotate-half** — PE-Core's interleaved rotary uses a strided `x[..., ::2]` that lowers to `GATHER_ND` (GPU-banned). Bake an even→odd channel permutation into the q/k weights (preserves q·k exactly) and apply the rotate-half form with constant cos/sin → clean `MUL`/`ADD`/`SLICE`/`CONCAT`. 3. **Attention-pool single-query attention → broadcast-multiply + reduce-sum** — the pooling query is a constant latent, so a batch-matmul there is `const @ non-const` (rejected at compile, and the reordered `const-RHS` form is mis-computed on device); expressing it as `(q·k).sum` + softmax + `(attn·v).sum` is exact and GPU-correct. 4. **Overflow-safe LayerNorm** — the delegate computes the LayerNorm variance reduction in **fp16 even for an fp32 graph**; deep-ViT "massive activations" (|x|~50+) make `sum((x-mean)²)` exceed fp16 max (65504), so the normalization is wrong and the error compounds with depth (output correlation collapses to ~0.28 over 12 blocks while *still reporting full GPU residency*). Scaling by 1/32 before squaring (undone after) keeps the running sum in range — mathematically identical to `nn.LayerNorm`. Verified **on a Pixel 8a GPU**: zero banned ops, zero >4D tensors, full residency, and TFLite(GPU)-vs-PyTorch output correlation = 1.0 (the on-device GPU result — not just the host CPU result — matches the reference). ## Training data & PII PE-Core was pretrained by Meta on a large-scale **web-crawled image–text dataset** (billions of image–caption pairs, CLIP-style contrastive objective). No new training was performed for this conversion — it is a weights-exact format change of the public `timm`/`facebook` checkpoint. Because the source data is web-scraped, it may incidentally contain people, faces, text, and other PII; no PII was deliberately collected, and this conversion adds none. Users deploying the encoder should apply their own content/PII filtering as appropriate. See the original [PE model card](https://huggingface.co/facebook/PE-Core-B16-224) and [paper](https://arxiv.org/abs/2504.13181) for full dataset details. ## License & attribution - **Apache-2.0** (original [PE-Core](https://huggingface.co/facebook/PE-Core-B16-224) / [timm checkpoint](https://huggingface.co/timm/vit_pe_core_base_patch16_224.fb)). - This is a format conversion; all credit to the original authors (Meta / FAIR).