litert-community/PE-Core-base-patch16-224

Perception Encoder (PE-Core-B16-224) — LiteRT (TFLite) GPU

On-device 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 (ViT-B/16, 94M params; original 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)

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. 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 and paper for full dataset details.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime	Backend	Graph on GPU	Latency
TFLite benchmark_model (TfLiteGpuDelegateV2)	GPU (OpenCL)	80 / 1102	did not run
TFLite benchmark_model	CPU (XNNPACK, 4 threads)	—	624.3 ms

Any on-device figure recorded when this model shipped came from a different runtime. It was taken through LiteRT's own CompiledModel accelerator (logcat reports it as LITERT_CL), which is the path the Kotlin sample app and the LiteRT API use, and it appears elsewhere on this card. The rows above are the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. The two are not comparable, so read the rows above as a reproducible floor rather than as this model's speed on LiteRT.

Note that the GPU does not take the whole graph here (80 of 1102); the remainder runs on the CPU and the split costs a per-partition round trip.

License & attribution

• Apache-2.0 (original PE-Core / timm checkpoint).
• This is a format conversion; all credit to the original authors (Meta / FAIR).