RTMW-m (Whole-Body) β€” LiteRT (on-device 133-keypoint pose, fully-GPU)

RTMW (mmpose, CSPNeXt + CSPNeXtPAFPN neck + RTMW/SimCC head) whole-body 2D pose, converted to LiteRT and running fully on the CompiledModel GPU (ML Drift) on Android. 133 COCO-WholeBody keypoints β€” 17 body + 6 feet + 68 face + 42 hands β€” for a single centered person.

RTMW β€” input | whole-body skeleton (on-device LiteRT GPU)

On-device (Pixel 8a, Tensor G3 β€” verified)

nodes on GPU 531 / 531 LITERT_CL (full residency)
inference ~6 ms (256Γ—192)
size 66 MB (fp16)
accuracy device-vs-PyTorch SimCC corr 0.999, keypoints within 0.2 px
image[1,3,256,192] (ImageNet 0-255) β†’[GPU: CSPNeXt + PAFPN + RTMW]β†’ simcc_x[1,133,384], simcc_y[1,133,512]

Minimal usage

Android (Kotlin, CompiledModel GPU)

val model = CompiledModel.create(context.assets, "rtmw_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(chw)              // [1,3,256,192] mmpose mean/std (0-255 RGB), NCHW
model.run(inputs, outputs)
val simccX = outputs[0].readFloat()    // [1,133,384]
val simccY = outputs[1].readFloat()    // [1,133,512]; keypoint = argmax / 2

Python (desktop verification)

MEAN = np.array([123.675, 116.28, 103.53], np.float32)
STD  = np.array([58.395, 57.12, 57.375], np.float32)
import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

img = Image.open("person.jpg").convert("RGB").resize((192, 256))  # centered subject crop
x = ((np.asarray(img, np.float32) - MEAN) / STD).transpose(2, 0, 1)[None]

it = Interpreter(model_path="rtmw_fp16.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], x); it.invoke()
od = it.get_output_details()
sx, sy = (it.get_tensor(o["index"])[0] for o in od)              # [133,384], [133,512]
if sx.shape[-1] != 384: sx, sy = sy, sx                          # identify by bin count
kx, ky = sx.argmax(-1) / 2.0, sy.argmax(-1) / 2.0                 # 133 keypoints, px in 192x256
for i, (a, b) in enumerate(zip(kx, ky)):
    print(f"kp{i}: ({a:.1f}, {b:.1f})")

How it converts (litert-torch)

The RTMPose-family re-authorings (all numerically exact) plus one extra for RTMW's neck/head:

  1. ScaleNorm (RMS) β†’ SafeRMSNorm β€” its input overflows fp16 (Ξ£xΒ²β‰ˆ3.6M > 65504) on Mali β†’ norm=∞ β†’ all-zero head; scale x down by S=64 before squaring.
  2. GAU act@act BMM β†’ broadcast-multiply + reduce-sum.
  3. nn.PixelShuffle β†’ depth-to-space ConvTranspose2d (ZeroStuffConvT2d) β€” the RTMW head's PixelShuffle upsample lowers to a 6D tensor (>4D, GPU-rejected); the fixed depth-to-space conv keeps it 4D and exact.

Result: banned ops NONE, all tensors ≀4D, tflite-vs-torch corr 1.0, device-vs-torch corr 0.999.

Preprocessing

Center-crop to 3:4, resize to 192Γ—256, ImageNet 0-255 normalize (mean [123.675, 116.28, 103.53], std [58.395, 57.12, 57.375]), NCHW. Top-down β€” one centered person. SimCC argmax (Γ· split=2) β†’ pixel.

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
LiteRT CompiledModel (LITERT_CL) GPU 531 / 531 ~6 ms
TFLite benchmark_model (TfLiteGpuDelegateV2) GPU (OpenCL) 531 / 531 38.5 ms
TFLite benchmark_model CPU (XNNPACK, 4 threads) β€” XNNPACK declined the graph

The two GPU rows are different runtimes, not a contradiction. The LITERT_CL figure is the one recorded when this model shipped, taken through LiteRT's own CompiledModel accelerator β€” the path the Kotlin sample app and the LiteRT API use. The TfLiteGpuDelegateV2 figure is the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. They agree on how much of the graph the GPU takes; they disagree on speed, and the classic delegate is the slower of the two here. Read the TfLiteGpuDelegateV2 row as a reproducible floor, not 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.

License

Apache-2.0. Upstream: open-mmlab/mmpose RTMW.

Downloads last month
31
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Collection including litert-community/RTMW-m-WholeBody-LiteRT