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.