wav2vec2-base-960h CTC β€” LiteRT (GPU)

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.

wav2vec2 CTC word onsets Real model output: char-CTC word onsets for J.F. Kennedy's 1961 inaugural address (U.S. National Archives recording, public domain).

Ships as two GPU graphs β€” the fused graph exceeds the Mali whole-graph shader-compile limit (a graph can be op-clean and still fail to compile when fused; each half compiles and runs fully delegated):

File Size Input Output API
w2v2_asr_frontend_fp16.tflite 9 MB waveform [1, 256000] features [1, 799, 768] CompiledModel GPU
w2v2_asr_head_fp16.tflite 180 MB features [1, 799, 768] CTC logits [1, 799, 32] CompiledModel GPU

Pipeline

16 kHz mono PCM in [-1, 1], zero-padded to the fixed 16 s window β†’ [GPU] conv frontend β†’ [GPU] 12-layer transformer + lm_head β†’ host greedy-CTC over the valid frames

  • Valid frames for n samples: run L=(L-k)//s+1 over the conv stack (10,5)(3,2)(3,2)(3,2)(3,2)(2,2)(2,2) β€” 50 Hz frames (16 s β†’ 799).
  • Blank id 0 (<pad>), | = word delimiter (tokens.txt, index-ordered).
  • Greedy char-CTC without an LM has the model's known spelling quirks on hard words (e.g. GRAVED/GRAVE); add a beam+LM host-side if you need the last WER point.

Minimal usage β€” Python

import numpy as np, torch, torchaudio
from ai_edge_litert.interpreter import Interpreter

wave, sr = torchaudio.load("speech.wav")             # 16 kHz mono, [-1,1]
x = torch.zeros(1, 256000); n = min(wave.shape[1], 256000)
x[0, :n] = wave[0, :n]

def run(path, inp):
    it = Interpreter(model_path=path); it.allocate_tensors()
    d = it.get_input_details()[0]
    it.set_tensor(d["index"], inp.astype(np.float32)); it.invoke()
    return it.get_tensor(it.get_output_details()[0]["index"])

feat = run("w2v2_asr_frontend_fp16.tflite", x.numpy())
logits = run("w2v2_asr_head_fp16.tflite", feat)[0]    # [799, 32]

L = n
for k, s in [(10,5),(3,2),(3,2),(3,2),(3,2),(2,2),(2,2)]:
    L = (L - k) // s + 1
tokens = open("tokens.txt").read().splitlines()
out, prev = [], -1
for i in logits[:L].argmax(-1):
    if i != prev and i != 0: out.append(tokens[int(i)])
    prev = i
print("".join(out).replace("|", " ").strip())

Minimal usage β€” Kotlin (Android)

val frontend = CompiledModel.create(frontendPath, CompiledModel.Options(Accelerator.GPU), null)
val head = CompiledModel.create(headPath, CompiledModel.Options(Accelerator.GPU), null)
val fIn = frontend.createInputBuffers(); val fOut = frontend.createOutputBuffers()
val hIn = head.createInputBuffers(); val hOut = head.createOutputBuffers()

fIn[0].writeFloat(pcm)                        // [-1,1] floats, zero-padded to 256000
frontend.run(fIn, fOut)
hIn[0].writeFloat(fOut[0].readFloat())        // features [1,799,768]
head.run(hIn, hOut)
val logits = hOut[0].readFloat()              // [799 * 32], readback syncs the GPU
// greedy CTC over the valid frames: argmax per frame, drop blanks (id 0) + repeats,
// map through tokens.txt, '|' -> space

On-device performance (Pixel 8a, CompiledModel GPU)

  • frontend 448 ms + head 391 ms per 16 s window (RTF β‰ˆ 0.05); GPU compile 0.7 s + 1.5 s.
  • Device logits vs desktop float reference: corr 0.9928 (valid region), per-frame argmax agreement 97.0 %; transcript matches the desktop reference.

Conversion notes

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.

Sources & license

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

Model tree for litert-community/wav2vec2-base-960h-LiteRT

Finetuned
(178)
this model

Paper for litert-community/wav2vec2-base-960h-LiteRT