litert-community/RT-DETRv2-S-LiteRT

RT-DETRv2-S — LiteRT (CompiledModel GPU)

RT-DETRv2 (Baidu, 2024 — PekingU/rtdetr_v2_r18vd) object detection, converted to LiteRT and running 100% on the CompiledModel GPU (ML Drift) on a phone, with no CPU/ONNX fallback.

RT-DETRv2 is a transformer detector (ResNet18-vd backbone + a hybrid AIFI/CCFM encoder + a plain deformable-attention DETR decoder). Off-the-shelf it is GPU-incompatible (deformable grid_sample → GATHER_ND, two-stage query selection → TOPK/GATHER). Here it is converted with litert-torch and split into two GPU graphs with a host step between them, so both transformer graphs run on the GPU.

Files

File	What it is	Size (fp16)
rtdetr_graphA_fp16.tflite	ResNet18-vd backbone + hybrid encoder + score head → enc_class[1,8400,80], memory_raw[1,8400,256]	33.8 MB
rtdetr_graphB_fp16.tflite	two-stage combine + plain decoder + heads → boxes[1,300,4] (cxcywh), logits[1,300,80]	7.7 MB
host_params.bin	host per-token tail weights (enc_output + enc_bbox_head), valid mask, anchors (fp32)	0.9 MB
coco_labels.txt	80 contiguous COCO class names (id 0–79)	—

How it runs (two-graph split)

image[1,3,640,640]
  →[GPU Graph A]→ enc_class, memory_raw
  →[host: top-300 by max class score; per-token tail on the 300 selected (fp32):
          target = enc_output(valid·memory_raw)   (Linear + LayerNorm)
          ref    = enc_bbox_head(target) + anchors (3-layer MLP)]
  →[GPU Graph B  (memory_raw, target, ref)]→ boxes[1,300,4], logits[1,300,80]
  →[host: sigmoid + threshold + cxcywh→xyxy + light NMS]→ detections

The two-stage query selection (TOPK/GATHER) has no GPU op, but the proposal grid is image-independent, so the model splits there. The per-token tail (enc_output + enc_bbox_head) runs on the host over the 300 selected tokens (exact, since per-token ops commute with the gather).

Why the per-token tail is on the host — a Mali 3D-sequence fan-out bug

Both graphs convert GPU-clean, but a naïve Graph A (emitting enc_class/enc_coord/output_memory/memory_raw together) silently produced wrong boxes on device — large objects vanished while small ones stayed perfect. A 3-D token tensor [1,N,256] (from conv.flatten(2).transpose(1,2)) that is both a graph output and consumed by another node — or that fans out to several consumers — gets clobbered on the longer branch (4-D conv-map outputs are fine). output_memory fed both heads; the 3-layer box head lost, so its reference-box deltas collapsed to ~0. Fix: Graph A emits only the two fp16-clean leaves (enc_class + memory_raw×2) and the per-token tail moves to the host.

Minimal usage

Android (Kotlin, CompiledModel GPU)

val ga = CompiledModel.create(context.assets, "rtdetr_graphA_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val gb = CompiledModel.create(context.assets, "rtdetr_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,640,640] RGB in [0,1], NCHW
ga.run(aIn, aOut)                        // -> enc_class[1,8400,80], memory_raw*2[1,8400,256]
// host step: /2 -> top-300 -> per-token tail (host_params.bin) -> target[1,300,256], ref[1,300,4]
// (resolve buffer slots by float size; full math in the Python below / litert-samples object_detection)
bIn[0].writeFloat(memory); bIn[1].writeFloat(target); bIn[2].writeFloat(ref)
gb.run(bIn, bOut)
val boxes = bOut[0].readFloat()          // [1,300,4] cxcywh in [0,1]
val logits = bOut[1].readFloat()         // [1,300,80] -> sigmoid + threshold + light NMS

Python (desktop verification)

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

NP_, NQ, NC, H = 8400, 300, 80, 256
img = Image.open("photo.jpg").convert("RGB").resize((640, 640))
x = (np.asarray(img, np.float32) / 255.0).transpose(2, 0, 1)[None]  # [1,3,640,640], [0,1] only

# host_params.bin (fp32 LE): enc_output W[256,256],b,gamma,beta · bbox-MLP W0,b0,W1,b1,W2[4,256],b2 · valid[8400] · anchors[8400,4]
p = np.fromfile("host_params.bin", np.float32); o = 0
def take(*s):
    global o; n = int(np.prod(s)); v = p[o:o+n].reshape(s); o += n; return v
eoW, eoB, eoG, eoBe = take(H, H), take(H), take(H), take(H)
W0, b0, W1, b1, W2, b2 = take(H, H), take(H), take(H, H), take(H), take(4, H), take(4)
valid, anchors = take(NP_), take(NP_, 4)

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("rtdetr_graphA_fp16.tflite", {(3, 640, 640): x})
enc_cls, mem = a[(NP_, NC)][0], a[(NP_, H)][0] / 2.0                # Graph A emits memory_raw*2 — undo

top = np.argsort(-enc_cls.max(-1))[:NQ]                             # top-300 by max class logit
t = (valid[top, None] * mem[top]) @ eoW.T + eoB                     # per-token tail: enc_output Linear...
t = (t - t.mean(-1, keepdims=True)) / np.sqrt(t.var(-1, keepdims=True) + 1e-5) * eoG + eoBe  # ...+ LayerNorm
h = np.maximum(t @ W0.T + b0, 0); h = np.maximum(h @ W1.T + b1, 0)
ref = h @ W2.T + b2 + anchors[top]                                  # enc_bbox_head MLP + anchors

b = run("rtdetr_graphB_fp16.tflite",
        {(NP_, H): mem[None], (NQ, H): t[None].astype(np.float32), (NQ, 4): ref[None].astype(np.float32)})
boxes, logits = b[(NQ, 4)][0], b[(NQ, NC)][0]                       # cxcywh in [0,1] / 80-way logits
labels = open("coco_labels.txt").read().splitlines()
score = 1 / (1 + np.exp(-logits.max(-1))); cls = logits.argmax(-1)
for q in np.where(score > 0.4)[0]:                                  # + light NMS (IoU 0.7) in a real app
    cx, cy, w, hh = boxes[q]
    print(f"{labels[cls[q]]:12s} {score[q]:.2f}  xyxy=({cx-w/2:.3f},{cy-hh/2:.3f},{cx+w/2:.3f},{cy+hh/2:.3f})")

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

Both graphs run 100% GPU-resident (LITERT_CL): Graph A fully delegated, Graph B 704/704. The device chain reproduces the PyTorch detections exactly — COCO val giraffe image 7/7, cats image (000000039769) 6/6, every box at IoU 0.98–1.00 with matching class and score.

End-to-end ~615 ms/frame on a Pixel 8a: Graph B's deformable decoder over RT-DETR's 8400 tokens / 80×80 levels is ~350 ms of GPU compute (the GATHER-free tent-matmul grid_sample turns an O(points) gather into an O(H·W) matmul). So this model is accurate and fully-GPU but not real-time on this device; it suits still-image / snapshot detection. (A real-time camera demo of the same family is RF-DETR Nano, whose single small deformable level runs at ~9 fps.)

Preprocessing / outputs

• Input: square resize to 640×640, RGB, [0,1] rescale only (no ImageNet normalization), NCHW.
• Output: Graph B boxes are cxcywh normalized to [0,1]; logits are 80-way (contiguous COCO id 0–79). Host applies sigmoid + score threshold + cxcywh→xyxy + light NMS.

Conversion notes

Converted with litert-torch (NCHW preserved — onnx2tf destroys ViT attention). Re-authoring (per-graph tflite-vs-torch correlation 1.0): deformable grid_sample → a GATHER/CAST-free tent-matmul, MSDeformAttn ≤4D, baked AIFI sine pos-embed, ResNet18-vd stem zero-pad maxpool (the -inf-pad maxpool lowers to a Mali-rejected PADV2), a down-scaled fp16-safe LayerNorm, and the 3D-fan-out fix above (emit clean leaves + host-side per-token tail).

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 lyuwenyu/RT-DETR.