litert-community/yolox-nano-litert

YOLOX-Nano — LiteRT (CompiledModel GPU)

Megvii YOLOX-Nano (COCO, Apache-2.0) re-authored to a GPU-native LiteRT .tflite via the official litert_torch path (no onnx2tf). FP16, 2.2 MB, input 416×416.

Verified on a Pixel 8a: the whole graph runs on the GPU delegate (full LITERT_CL residency, zero CPU fallback) and the GPU output matches the CPU/PyTorch reference (corr ≥ 0.999).

Why this is GPU-clean

YOLOX is a pure CNN, but its Focus stem (stride-2 space-to-depth slicing) lowers to GATHER_ND, which the GPU delegate rejects. Here the Focus + its following 3×3 conv are folded into a single, numerically-exact 6×6 stride-2 conv, so the graph has zero GATHER/GATHER_ND/ TopK/Cast ops and no >4D tensors. Activations (SiLU) lower to LOGISTIC+MUL.

I/O

• Input images [1, 416, 416, 3] NHWC, BGR, 0–255, no normalization (YOLOX letterbox: uniform-scale to fit, pad bottom/right with gray 114).
• Output [1, 3549, 85] raw heads, anchor-major. `85 = 4 box (cx,cy,w,h, grid units) + 1 obj
  • 80 class`. obj/class are already sigmoid'd; boxes are not decoded.

Host-side decode (kept out of the graph for GPU-cleanliness)

For anchor i at grid (gx,gy) with stride ∈ {8,16,32}: cx=(raw_cx+gx)*stride, cy=(raw_cy+gy)*stride, w=exp(raw_w)*stride, h=exp(raw_h)*stride; score = obj * max_class; then per-class NMS. Divide boxes by the letterbox ratio to map back. Reference Kotlin + Python decode in the sample below.

Minimal usage

Android (Kotlin, CompiledModel GPU)

val model = CompiledModel.create(context.assets, "yolox_nano.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(nhwc)             // [1,416,416,3] BGR 0-255, letterbox pad 114
model.run(inputs, outputs)
val raw = outputs[0].readFloat()       // [1,3549,85] -> decode + NMS on host (see Python)

Python (desktop verification)

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

SIZE = 416
img = Image.open("photo.jpg").convert("RGB")
r = min(SIZE / img.width, SIZE / img.height)
w, h = round(img.width * r), round(img.height * r)
canvas = np.full((SIZE, SIZE, 3), 114, np.float32)                # letterbox, gray 114
canvas[:h, :w] = np.asarray(img.resize((w, h)), np.float32)
x = np.ascontiguousarray(canvas[..., ::-1])[None]                 # RGB -> BGR, 0-255, NHWC

it = Interpreter(model_path="yolox_nano.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], x); it.invoke()
out = it.get_tensor(it.get_output_details()[0]["index"])[0]       # [3549,85]

grids, strides = [], []                                           # anchors = grid cells, s 8/16/32
for s in (8, 16, 32):
    n = SIZE // s
    gy, gx = np.mgrid[:n, :n]
    grids.append(np.stack([gx, gy], -1).reshape(-1, 2)); strides.append(np.full((n * n, 1), s))
g = np.concatenate(grids).astype(np.float32); sv = np.concatenate(strides).astype(np.float32)
xy = (out[:, :2] + g) * sv; wh = np.exp(out[:, 2:4]) * sv         # boxes in 416-space
score = out[:, 4:5] * out[:, 5:]                                  # obj x class (already sigmoid)
cls, conf = score.argmax(1), score.max(1)
for i in np.where(conf > 0.35)[0]:                                # + per-class NMS in practice
    x1, y1 = (xy[i] - wh[i] / 2) / r; x2, y2 = (xy[i] + wh[i] / 2) / r
    print(f"coco class {cls[i]}  {conf[i]:.2f}  [{x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f}]")

Performance

COCO val2017 AP 25.8 (FP32 reference). Real-time on Pixel 8a GPU.

Training data & PII

Trained by Megvii on COCO 2017 (train2017), a public academic object-detection dataset (Creative Commons). COCO images contain people as one of the 80 object categories; no names, identities, or other personal attributes are modeled or output — the model emits only class id + box. No additional or private data was used. Weights are the official Megvii release; only the op graph was re-authored for GPU (weights unchanged).

Sample app + conversion script

Android sample (CompiledModel GPU, Kotlin decode + NMS) and the litert_torch conversion script: https://github.com/google-ai-edge/litert-samples (compiled_model_api/object_detection)

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)	GPU (OpenCL)	482 / 482	24.0 ms
TFLite benchmark_model	CPU (XNNPACK, 4 threads)	—	XNNPACK declined the graph

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.

XNNPACK declines these fp16 graphs — it reports failed to delegate DEPTHWISE_CONV_2D and then fails to allocate tensors — so there is no usable CPU number. Disabling XNNPACK falls back to reference kernels, which measured about 20× slower than the GPU on models of this size and would not represent CPU inference anyone would ship.