litert-community/RF-DETR-Nano-LiteRT

RF-DETR Nano — LiteRT (CompiledModel GPU)

RF-DETR (Roboflow 2025, an LW-DETR derivative) object detection, converted to LiteRT and running 100% on the CompiledModel GPU (ML Drift) on a phone — the first transformer/DETR detector to ride the LiteRT GPU API with no CPU/ONNX fallback.

RF-DETR is a transformer detector (windowed DINOv2-S backbone + deformable-attention DETR decoder). Off-the-shelf it is GPU-incompatible (deformable grid_sample → GATHER_ND, windowed attention → 5D/6D tensors, two-stage query selection → TOPK/GATHER). Here it is converted with litert-torch and split into two GPU graphs with a tiny host step between them, so the whole detector runs on the GPU.

Files

File	What it is	Size (fp16)
rfdetr_graphA_fp16.tflite	backbone + encoder + proposal heads → enc_class[1,576,91], enc_coord[1,576,4], memory[1,576,256]	48.6 MB
rfdetr_graphB_fp16.tflite	two-stage combine + decoder + heads → boxes[1,300,4] (cxcywh), logits[1,300,91]	7.6 MB

How it runs (two-graph split)

image[1,3,384,384]
  →[GPU Graph A]→ enc_class, enc_coord, memory
  →[host: top-300 by max class score → gather coords]→ refpoint_ts[1,300,4]
  →[GPU Graph B  (memory, refpoint_ts)]→ boxes[1,300,4], logits[1,300,91]
  →[host: sigmoid + threshold + cxcywh→xyxy + per-class NMS]→ detections

The two-stage query selection (TOPK/GATHER) has no GPU op, but the proposal grid is image-independent, so the model splits at exactly that point — the standard two-stage-DETR edge split. Both graphs are 100% GPU-resident.

Minimal usage

Android (Kotlin, CompiledModel GPU)

val ga = CompiledModel.create(context.assets, "rfdetr_graphA_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val gb = CompiledModel.create(context.assets, "rfdetr_graphB_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val aIn = ga.createInputBuffers(); val aOut = ga.createOutputBuffers()
val bIn = gb.createInputBuffers(); val bOut = gb.createOutputBuffers()
aIn[0].writeFloat(chw)               // [1,3,384,384] RGB, ImageNet mean/std, NCHW
ga.run(aIn, aOut)                    // -> enc_class[1,576,91], enc_coord[1,576,4], memory[1,576,256]
// host step: top-300 by max class logit -> gather enc_coord -> refpoint_ts[1,300,4]
// (resolve buffer slots by float size; see the Python below / litert-samples object_detection sample)
bIn[0].writeFloat(memory); bIn[1].writeFloat(refpointTs)
gb.run(bIn, bOut)
val boxes = bOut[0].readFloat()      // [1,300,4] cxcywh in [0,1]
val logits = bOut[1].readFloat()     // [1,300,91] -> sigmoid + threshold + per-class NMS

Python (desktop verification)

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

NP_, NQ, NC, H = 576, 300, 91, 256
img = Image.open("photo.jpg").convert("RGB").resize((384, 384))
x = np.asarray(img, np.float32) / 255.0
x = ((x - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]).astype(np.float32).transpose(2, 0, 1)[None]

def run(path, feeds):  # feed/fetch tensors by shape (converter slot order is arbitrary)
    it = Interpreter(model_path=path); it.allocate_tensors()
    for d in it.get_input_details(): it.set_tensor(d["index"], feeds[tuple(d["shape"][1:])])
    it.invoke(); return {tuple(d["shape"][1:]): it.get_tensor(d["index"]) for d in it.get_output_details()}

a = run("rfdetr_graphA_fp16.tflite", {(3, 384, 384): x})
enc_cls, enc_coord, mem = a[(NP_, NC)][0], a[(NP_, 4)][0], a[(NP_, H)][0]

top = np.argsort(-enc_cls.max(-1))[:NQ]                     # top-300 by max class logit
ts = enc_coord[top]                                         # gather -> refpoint_ts [300,4]

b = run("rfdetr_graphB_fp16.tflite", {(NP_, H): mem[None], (NQ, 4): ts[None]})
boxes, logits = b[(NQ, 4)][0], b[(NQ, NC)][0]               # cxcywh in [0,1] / 91-way (index = COCO category id)
score = 1 / (1 + np.exp(-logits.max(-1))); cls = logits.argmax(-1)
for q in np.where((score > 0.45) & (cls > 0))[0]:           # id 0 unused; + per-class NMS (IoU 0.6) in a real app
    cx, cy, w, h = boxes[q]
    print(f"COCO id {cls[q]:2d}  {score[q]:.2f}  xyxy=({cx-w/2:.3f},{cy-h/2:.3f},{cx+w/2:.3f},{cy+h/2:.3f})")

On-device (Pixel 8a, Tensor G3 — verified)

graph	nodes on GPU	time
Graph A	1381/1381 LITERT_CL	~22 ms
Graph B	404/404 LITERT_CL	~5 ms

Full pipeline ≈ 27 ms (model) / ~100 ms end-to-end incl. host pre/post-processing. On a real image the device chain reproduces the PyTorch detections at IoU 0.98–0.99 with matching class and score.

Preprocessing / outputs

• Input: square resize to 384×384, RGB, ImageNet mean/std ([0.485,0.456,0.406]/[0.229,0.224,0.225]), NCHW.
• Output: Graph B boxes are cxcywh normalized to [0,1]; logits are 91-way (index = COCO category id). Host applies sigmoid + score threshold + cxcywh→xyxy + per-class NMS.

Conversion notes

Converted with litert-torch (NCHW preserved — onnx2tf destroys ViT attention). Re-authoring (per-graph tflite-vs-torch correlation 1.0): windowed DINOv2 backbone (6D window-partition → ≤4D, SDPA → manual attention), deformable grid_sample → a GATHER/CAST-free tent-matmul, MSDeformAttn ≤4D, baked sine pos-embed, and a down-scaled fp16-safe LayerNorm in the projector and decoder (the Mali delegate computes in fp16, and those LayerNorm channel-sums otherwise overflow). The two-stage topk/gather runs on the host between the two graphs.

A runnable Android sample (CompiledModel GPU) and the conversion scripts are in the official ai-edge-litert/litert-samples object_detection example.

License

Apache-2.0, inherited from roboflow/rf-detr.