litert-community/RF-DETR-Seg-Nano-LiteRT

RF-DETR-Seg Nano — LiteRT (CompiledModel GPU)

Live output on a Pixel 8a (photo: Pexels, free license).

RF-DETR-Seg (Roboflow, rf-detr 1.9.3) instance segmentation, converted to LiteRT and running 100% on the CompiledModel GPU (ML Drift) on a phone — a DETR-family segmenter (DINOv2-S/12 backbone + deformable-attention decoder + ConvNeXt-style mask head, 33.6M params, COCO seg AP₅₀ 63.0) with no CPU/ONNX fallback.

Off-the-shelf it is GPU-incompatible (deformable grid_sample → GATHER_ND, two-stage query selection → TOPK/GATHER, SDPA → rank-3 batched matmuls, and several large baked constants that the GPU delegate executes incorrectly). Here it is converted with litert-torch, re-authored op-by-op, and split into two GPU graphs with a tiny host step between them.

Files

File	What it is	Size
rfdetrseg_graphA_fp16.tflite	backbone + encoder + proposal heads → enc_class[1,676,91], enc_delta[1,676,4], memory×2[1,676,256]	47.0 MB
rfdetrseg_graphB_fp16.tflite	decoder + box/class heads + mask head → boxes[1,100,4] (cxcywh), logits[1,100,91], masks[1,100,78,78]	14.8 MB
clspos.bin	host-fed constant cls_token + pos_embed[:, :1], float32 [1,1,384]	1.5 KB
pospatch.bin	host-fed constant patch pos-embed, float32 [1,676,384]	1.0 MB
query_feat.bin	host-fed constant decoder query embedding, float32 [1,100,256]	100 KB
refpoint_embed.bin	learned reference points for the host reparam, float32 [1,100,4]	1.6 KB

How it runs (two-graph split + host-fed constants)

image[1,3,312,312] + clspos[1,1,384] + pospatch[1,676,384]
  →[GPU Graph A]→ enc_class[1,676,91], enc_delta[1,676,4], memory×2[1,676,256]
  →[host: ÷2 → proposal-grid combine → top-100 by max class score
          → gather → reparam with refpoint_embed]→ refpoint[1,100,4]
  →[GPU Graph B  (memory, refpoint, query_feat)]→ boxes, logits, masks[1,100,78,78]
  →[host: sigmoid + threshold + per-class NMS]→ instances (mask inside = logit > 0)

The proposal grid is image-independent (26×26, cxcy = (grid+0.5)/26, wh = 0.05), so the host step is pure elementwise math plus a topk.

Why three constants are graph inputs: the ML Drift GPU delegate silently mis-executes compute chains that consume large baked-constant tensors (fp32 and fp16 return identical wrong numbers — not a precision issue). The cls+pos embedding, the patch pos-embed and the decoder query embedding are therefore fed at runtime from the .bin files above, and the reparam combine (which would consume the baked refpoint_embed) runs on the host. Graph A also emits memory×2 because a [1,N,C] tensor that is both consumed and output comes back zeroed on the delegate — halve it on the host.

Minimal usage

Android (Kotlin, CompiledModel GPU)

val env = Environment.create()                      // ONE shared env for both graphs
val ga = CompiledModel.create(pathA, CompiledModel.Options(Accelerator.GPU), env)
val gb = CompiledModel.create(pathB, CompiledModel.Options(Accelerator.GPU), env)
val aIn = ga.createInputBuffers(); val aOut = ga.createOutputBuffers()
val bIn = gb.createInputBuffers(); val bOut = gb.createOutputBuffers()
// resolve slots by float size (converter order is arbitrary); write the .bin constants once:
aIn[clsposSlot].writeFloat(clspos); aIn[pospatchSlot].writeFloat(pospatch)
bIn[qfSlot].writeFloat(queryFeat)
aIn[imageSlot].writeFloat(chw)                      // [1,3,312,312] RGB, ImageNet mean/std
ga.run(aIn, aOut)
// host: memory = memory2 * 0.5; proposal combine + top-100 + gather + reparam -> refpoint[1,100,4]
bIn[memSlot].writeFloat(memory); bIn[refSlot].writeFloat(refpoint)
gb.run(bIn, bOut)
val boxes = bOut[boxSlot].readFloat()               // [1,100,4] cxcywh in [0,1]
val logits = bOut[logitSlot].readFloat()            // [1,100,91] -> sigmoid + threshold + NMS
val masks = bOut[maskSlot].readFloat()              // [1,100,78,78] full-image logits, inside = > 0

Python (desktop verification, CompiledModel API)

import numpy as np
from PIL import Image
from ai_edge_litert.compiled_model import CompiledModel

R, NP_, NQ, NC, H, M, G = 312, 676, 100, 91, 256, 78, 26
img = Image.open("photo.jpg").convert("RGB").resize((R, R))
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 buffers by float count (converter slot order is arbitrary)
    m = CompiledModel.from_file(path)
    ins, outs = m.create_input_buffers(0), m.create_output_buffers(0)
    for i in range(len(ins)):
        n = m.get_input_buffer_requirements(i)["buffer_size"] // 4
        ins[i].write(np.ascontiguousarray(feeds[n].ravel(), np.float32))
    m.run_by_index(0, ins, outs)
    return {m.get_output_buffer_requirements(j)["buffer_size"] // 4:
            outs[j].read(m.get_output_buffer_requirements(j)["buffer_size"] // 4, np.float32)
            for j in range(len(outs))}

clspos = np.fromfile("clspos.bin", np.float32)             # [1,1,384]
pospatch = np.fromfile("pospatch.bin", np.float32)         # [1,676,384]
rp = np.fromfile("refpoint_embed.bin", np.float32).reshape(NQ, 4)
qf = np.fromfile("query_feat.bin", np.float32)             # [1,100,256]

a = run("rfdetrseg_graphA_fp16.tflite", {x.size: x, clspos.size: clspos, pospatch.size: pospatch})
enc_cls = a[NP_ * NC].reshape(NP_, NC)
delta = a[NP_ * 4].reshape(NP_, 4)
mem = a[NP_ * H] * 0.5                                     # graph outputs memory*2

gy, gx = np.mgrid[0:G, 0:G]                                # proposal grid (image-independent)
prop = np.stack([(gx + .5) / G, (gy + .5) / G], -1).reshape(NP_, 2)
cxcy = delta[:, :2] * 0.05 + prop
wh = np.exp(delta[:, 2:]) * 0.05
top = np.argsort(-enc_cls.max(-1))[:NQ]                    # top-100 by max class logit
ts = np.concatenate([cxcy, wh], -1)[top]
ref = np.concatenate([rp[:, :2] * ts[:, 2:] + ts[:, :2], np.exp(rp[:, 2:]) * ts[:, 2:]], -1)

b = run("rfdetrseg_graphB_fp16.tflite", {mem.size: mem, ref.size: ref, qf.size: qf})
boxes = b[NQ * 4].reshape(NQ, 4)                           # cxcywh in [0,1]
logits = b[NQ * NC].reshape(NQ, NC)                        # 91-way (index = COCO category id)
masks = b[NQ * M * M].reshape(NQ, M, M)                    # full-image raw logits, inside = > 0

score = 1 / (1 + np.exp(-logits.max(-1))); cls = logits.argmax(-1)
for q in np.where((score > 0.5) & (cls > 0))[0]:           # + per-class NMS (IoU 0.6) in a real app
    print(f"COCO id {cls[q]:2d}  {score[q]:.2f}  cxcywh={np.round(boxes[q], 3)}  mask px={int((masks[q] > 0).sum())}")

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

graph	nodes on GPU	time
Graph A	1293/1293 LITERT_CL, 1 partition	17.5 ms
Graph B	884/884 LITERT_CL, 1 partition	9.1 ms

Real-image end-to-end (device chain vs the official PyTorch RFDETRSegNano.predict, threshold 0.5): every detection matches with box IoU ≥ 0.99, mask IoU ≥ 0.995 and identical classes on the test images (4/4 and 2/2 detections).

Street scene on the Pixel 8a: 10 instances (photo: Pexels, free license).

Preprocessing / outputs

• Input: square resize to 312×312, RGB, ImageNet mean/std ([0.485,0.456,0.406] / [0.229,0.224,0.225]), NCHW, plus the two constant embedding inputs (clspos.bin, pospatch.bin).
• Output: boxes are cxcywh normalized to [0,1]; logits are 91-way (index = COCO category id, id 0 unused); masks are per-query full-image 78×78 raw logits (sigmoid > 0.5 ⇔ logit > 0), upsample bilinearly to the frame.

Conversion notes

Converted from the PyTorch checkpoint (rfdetr 1.9.3, RFDETRSegNano) with litert-torch (NCHW preserved) + fp16 weights via ai-edge-quantizer float-casting. GPU re-authoring, all numerically exact:

• SDPA / nn.MultiheadAttention → manual rank-4 attention (the delegate mis-executes rank-3 batched matmuls).
• Deformable grid_sample → GATHER/CAST-free tent-matmul bilinear sampler (rank-4 BMM).
• SafeLayerNorm: adaptive per-row down-scale that never reconstructs the large variance — fp16-safe at any magnitude; channels-first sites use a 3D [B,HW,C] detour.
• tanh-GELU (no ERF lowering); sine pos-embed dim_t baked, interleave via reshape; seg einsum → rank-4 matmul.
• LayerScale folded into the preceding Linear, and the cls/pos/query embeddings host-fed (the baked-constant execution bug above).

Original project: roboflow/rf-detr (RF-DETR-Seg Nano, tag 1.9.3) — Apache-2.0. The rfdetr package and the Apache-designated checkpoints (including RF-DETR-Seg-N) are Apache 2.0; only the separate rfdetr_plus components are under PML 1.0 (not used here).