litert-community/RAM-Plus-LiteRT

RAM++ (Recognize Anything Plus) — LiteRT on-device image tagging

RAM++ (Apache-2.0) re-authored for LiteRT: give it a photo, get the tags it recognizes from a 4,585-tag open vocabulary — per-tag sigmoid, no fixed class head. Four graphs — the Swin-L encoder stages 0-2 and the Query2Label tag head run on the CompiledModel GPU; the last Swin stage and the 479 MB frozen tag bank run on CPU (the deep Swin block fp16-miscomputes on the Mali delegate — see below).

Verified on a Pixel 8a: Swin 0-2 GPU (corr 0.998) + stage-3/reweight CPU (exact) + tag head GPU (corr 0.9987, ~270 ms). Sample photo (a dog on a couch) → 14 tags in ~2 s, all correct: dog · couch · living room · sit · carpet · picture frame · plant · armchair · lamp · pillow ….

Files

file	graph	in → out	delegate
ram_swin_s012_fp16.tflite	Swin stages 0-2	image [1,3,384,384] → feat [1,144,1536]	GPU
ram_stage3_tail_fp16.tflite	Swin stage 3 + norm + proj	feat → image_embeds [1,145,512]	CPU
ram_reweight_fp16.tflite	multi-grained reweight	cls [1,512] → tag queries [1,4585,768]	CPU
ram_taghead_fp16.tflite	Query2Label tag head	queries + image_embeds → logits [1,4585]	GPU
ram_tag_list.txt, ram_tag_threshold.bin	host assets (4585 tags + per-class thresholds)	—	—

Pipeline

image →[ImageNet norm]→ [GPU Swin 0-2]→ feat →[CPU Swin-3 + norm + proj]→ image_embeds[1,145,512]
       token0 = cls →[CPU reweight over the 4585×51 tag bank]→ queries[1,4585,768]
       (queries, image_embeds) →[GPU Q2L tag head]→ logits →[sigmoid + per-class threshold]→ tags

Why the GPU/CPU split — a Mali fp16 finding

The Swin-L encoder is fully GPU-convertible, but its last stage miscomputes in fp16 on the Mali delegate. Bisecting the four stages on-device: stage 0 = 0.9999, stage 1 = 0.9999, stage 2 = 0.9983, stage 3 = 0.709. It is not head_dim (stage 2 shares head_dim 32) and not overflow (every stage-3 value < 848 ≪ fp16 max 65504; a round-to-fp16-between-ops simulation reproduces fp32 at corr 0.99999997) — it is Mali's fp16 matmul accumulation in the deep, high-magnitude blocks (the residual stream grows to absmax 847; the 6144-wide fc2 and 48-head attention accumulate in fp16). Those 2 blocks run on CPU; everything else stays on GPU. The reweight bakes the tag bank once as fp16 (229 MB, not 686 MB).

Minimal usage (Python)

import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

def run(path, x, *ins):                                  # single-input or size-matched multi-input
    it = Interpreter(model_path=path); it.allocate_tensors()
    ind = it.get_input_details()
    if not ins:
        it.set_tensor(ind[0]["index"], x)
    else:
        for d in ind:
            n = int(np.prod(d["shape"]))
            it.set_tensor(d["index"], x if n == x.size else ins[0])
    it.invoke()
    return it.get_tensor(it.get_output_details()[0]["index"])

# preprocess (ImageNet)
img = Image.open("photo.jpg").convert("RGB").resize((384, 384))
a = np.asarray(img, np.float32) / 255.0
a = (a - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
x = a.transpose(2, 0, 1)[None].astype(np.float32)        # [1,3,384,384]

feat = run("ram_swin_s012_fp16.tflite", x)               # [1,144,1536]
iemb = run("ram_stage3_tail_fp16.tflite", feat)          # [1,145,512]
cls = iemb[:, 0, :]                                       # [1,512]
queries = run("ram_reweight_fp16.tflite", cls)           # [1,4585,768]
logits = run("ram_taghead_fp16.tflite", queries, iemb)   # [1,4585]

probs = 1 / (1 + np.exp(-logits[0]))
thr = np.fromfile("ram_tag_threshold.bin", np.float32)
tags = [t for t in open("ram_tag_list.txt").read().splitlines()]
print([tags[i] for i in np.where(probs > thr)[0]])

Minimal usage (Kotlin, LiteRT CompiledModel)

val g1 = CompiledModel.create("ram_swin_s012_fp16.tflite", CompiledModel.Options(Accelerator.GPU), null)
val c2 = CompiledModel.create("ram_stage3_tail_fp16.tflite", CompiledModel.Options(Accelerator.CPU), null)
val rw = CompiledModel.create("ram_reweight_fp16.tflite", CompiledModel.Options(Accelerator.CPU), null)
val th = CompiledModel.create("ram_taghead_fp16.tflite", CompiledModel.Options(Accelerator.GPU), null)

g1In[0].writeFloat(preprocess(bitmap)); g1.run(g1In, g1Out)           // -> feat[1,144,1536]
c2In[0].writeFloat(g1Out[0].readFloat()); c2.run(c2In, c2Out)        // -> image_embeds[1,145,512]
val iemb = c2Out[0].readFloat(); val cls = iemb.copyOfRange(0, 512)
rwIn[0].writeFloat(cls); rw.run(rwIn, rwOut)                          // -> queries[1,4585,768]
val q = rwOut[0].readFloat()
for (b in thIn) { val n = b.readFloat().size; b.writeFloat(if (n == q.size) q else iemb) }
th.run(thIn, thOut)                                                  // -> logits[1,4585]
// sigmoid(logits[i]) > threshold[i] -> tag[i]

A complete Android sample (image pick → tags) is in google-ai-edge/litert-samples.

Upstream

xinyu1205/recognize-anything · xinyu1205/recognize-anything-plus-model (Apache-2.0). Paper: Open-Set Image Tagging with Multi-Grained Text Supervision.

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) — ram_taghead_fp16.tflite	GPU (OpenCL)	30 / 139	3499.7 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) — ram_reweight_fp16.tflite	GPU (OpenCL)	9 / 18	612.0 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) — ram_swin_s012_fp16.tflite	GPU (OpenCL)	105 / 2129	did not run
TFLite benchmark_model (TfLiteGpuDelegateV2) — ram_stage3_tail_fp16.tflite	GPU (OpenCL)	41 / 178	523.5 ms
TFLite benchmark_model — ram_taghead_fp16.tflite	CPU (XNNPACK, 4 threads)	—	1550.1 ms
TFLite benchmark_model — ram_reweight_fp16.tflite	CPU (XNNPACK, 4 threads)	—	312.8 ms
TFLite benchmark_model — ram_swin_s012_fp16.tflite	CPU (XNNPACK, 4 threads)	—	2786.9 ms
TFLite benchmark_model — ram_stage3_tail_fp16.tflite	CPU (XNNPACK, 4 threads)	—	217.0 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.

On this delegate the CPU is the faster choice for ram_taghead_fp16.tflite (1550.1 ms on CPU against 3499.7 ms on GPU), ram_reweight_fp16.tflite (312.8 ms on CPU against 612.0 ms on GPU), ram_stage3_tail_fp16.tflite (217.0 ms on CPU against 523.5 ms on GPU) — worth knowing before you reach for the GPU on a mid-range phone.

Note that the GPU does not take the whole graph here (30 / 139 in ram_taghead_fp16.tflite, 9 / 18 in ram_reweight_fp16.tflite, 105 / 2129 in ram_swin_s012_fp16.tflite, 41 / 178 in ram_stage3_tail_fp16.tflite); the remainder runs on the CPU and the split costs a per-partition round trip.