English speech recognition with
wav2vec2-base-960h
running
fully on the LiteRT CompiledModel GPU (ML Drift) — and with
zero FFT anywhere:
the raw 16 kHz waveform goes straight into the 1D-conv feature extractor, so there is no mel/fbank
step even on the host. Character-level CTC (29 chars + specials), greedy decode, no language model.
1import numpy as np, torch, torchaudio
2from ai_edge_litert.interpreter import Interpreter
3
4wave, sr = torchaudio.load("speech.wav") # 16 kHz mono, [-1,1]
5x = torch.zeros(1, 256000); n = min(wave.shape[1], 256000)
6x[0, :n] = wave[0, :n]
7
8def run(path, inp):
9 it = Interpreter(model_path=path); it.allocate_tensors()
10 d = it.get_input_details()[0]
11 it.set_tensor(d["index"], inp.astype(np.float32)); it.invoke()
12 return it.get_tensor(it.get_output_details()[0]["index"])
13
14feat = run("w2v2_asr_frontend_fp16.tflite", x.numpy())
15logits = run("w2v2_asr_head_fp16.tflite", feat)[0] # [799, 32]
16
17L = n
18for k, s in [(10,5),(3,2),(3,2),(3,2),(3,2),(2,2),(2,2)]:
19 L = (L - k) // s + 1
20tokens = open("tokens.txt").read().splitlines()
21out, prev = [], -1
22for i in logits[:L].argmax(-1):
23 if i != prev and i != 0: out.append(tokens[int(i)])
24 prev = i
25print("".join(out).replace("|", " ").strip())
1val frontend = CompiledModel.create(frontendPath, CompiledModel.Options(Accelerator.GPU), null)
2val head = CompiledModel.create(headPath, CompiledModel.Options(Accelerator.GPU), null)
3val fIn = frontend.createInputBuffers(); val fOut = frontend.createOutputBuffers()
4val hIn = head.createInputBuffers(); val hOut = head.createOutputBuffers()
5
6fIn[0].writeFloat(pcm) // [-1,1] floats, zero-padded to 256000
7frontend.run(fIn, fOut)
8hIn[0].writeFloat(fOut[0].readFloat()) // features [1,799,768]
9head.run(hIn, hOut)
10val logits = hOut[0].readFloat() // [799 * 32], readback syncs the GPU
11// greedy CTC over the valid frames: argmax per frame, drop blanks (id 0) + repeats,
12// map through tokens.txt, '|' -> space
Converted with litert-torch, numerically exact (tflite vs PyTorch: corr 1.000000):
GELU → tanh-GELU; frontend GroupNorm → 4D-reshape group-norm (avoids GATHER_ND);
pos_conv weight-norm folded to a static weight; the all-valid bidirectional attention mask
removed (fixed window → plain SDPA). The CTC head is a plain Linear — logits come out raw.
-
w2v2_asr_frontend_fp16.tflite — the NPU compiles this graph and then fails to run it: LiteRtException: Failed to invoke the compiled model. The GPU row below is the only S26 figure for it. A clean compile is not evidence that a model runs.
-
w2v2_asr_head_fp16.tflite — the NPU is 1.37x faster than the GPU (76.18 ms against 104.1 ms) and loads 5.84x faster (226 ms against 1322 ms).
The NPU rows here ran artifacts compiled ahead of time for SM8850 with QAIRT 2.47.0; the GPU rows ran the published files as they are. LiteRT can also compile for the NPU on the device at first load, which is what lets you ship the published file unchanged — that path and the ten runtime libraries it needs are in the
NPU recipe, and we did not measure it here. GPU wiring is in the
GPU recipe.
Measured on a Raspberry Pi 5 Model B Rev 1.1 (8 GB, Raspberry Pi OS 64-bit) with the LiteRT
benchmark_model tool from litert-cli-nightly 0.2.0.dev20260805: CPU inference (XNNPACK, 4 threads), 3 invocations per file of 10 warm-up plus 50 timed runs (the tool caps a phase at 150 s, so very slow graphs run fewer — the Runs column is the actual timed total). The latency is the median across invocations; the spread is the min–max over all timed runs. No thermal throttling occurred during these runs (
vcgencmd get_throttled stayed
0x0).