--- license: apache-2.0 library_name: litert pipeline_tag: object-detection tags: - object-detection - yolox - litert - tflite - on-device - gpu --- # YOLOX-Tiny — LiteRT (CompiledModel GPU) ![YOLOX-Tiny — on-device detections (Pixel 8a, LiteRT CompiledModel GPU)](samples/sample.png) Megvii **YOLOX-Tiny** (COCO, Apache-2.0) re-authored to a **GPU-native** LiteRT `.tflite` via the official **litert_torch** path (no onnx2tf). FP16, **10.4 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)** ```kotlin val model = CompiledModel.create(context.assets, "yolox_tiny.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)** ```python 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_tiny.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 **32.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)