litert-community/yolox-m-litert

YOLOX-M — LiteRT (CompiledModel GPU)

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

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, 640, 640, 3] NHWC, BGR, 0–255, no normalization (YOLOX letterbox: uniform-scale to fit, pad bottom/right with gray 114).
• Output [1, 8400, 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_m.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(nhwc)             // [1,640,640,3] BGR 0-255, letterbox pad 114
model.run(inputs, outputs)
val raw = outputs[0].readFloat()       // [1,8400,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 = 640
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_m.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]       # [8400,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 640-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 46.9 (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)