Lab 8.2 — On-Device Keyword / Tone Classifier

Course 2 syllabus · Bonus Module 8 · Prev: « Lab 8.1 · Next: Lab 8.3 »

Goal

Train a small convolutional neural network on the log-mel features from Lab 8.1 to classify a handful of spoken keywords (or DTMF tone pairs), then deploy it on-device — int8 with TensorRT on the Jetson and tflite / onnxruntime on the Pi 5 — and measure real latency and accuracy on the board, not on the host. The skill built is the full edge-ML deployment loop that DSP/firmware roles increasingly ask for: featurize → train on the host → quantize → run under a real-time budget on constrained hardware. The reference point is the classical Goertzel tone detector of Lab 6.5: for a pure-tone problem a five-line Goertzel beats a neural net on every axis, and knowing when the learned approach earns its keep is the point.

Equipment & parts

  • Host Mac for training (PyTorch), plus export to ONNX.
  • Jetson Orin Nano — TensorRT (trtexec / the Python API) for the int8 engine.
  • Raspberry Pi 5onnxruntime and/or tflite-runtime.
  • The audio input from Lab 8.1 (USB mic / I²S / ADS1115).
  • Optional known-tone source: MCP4725/STM32 DAC to synthesize repeatable DTMF pairs for a clean, labeled dataset and for head-to-head timing against Goertzel.

Safety & don’t-break-it

  • Same analog-input rules as Lab 8.1: keep the ADC/line input within rails, bias to mid-rail, common ground.
  • Don’t trust host accuracy as on-device accuracy. Quantization (fp32→int8) changes the numbers; always re-measure accuracy on the board with the board’s own microphone. A model that is 98% on the Mac can drop several points on int8 with a different mic frequency response.
  • Thermals again: a throttled Jetson inflates latency. Fix the clocks (jetson_clocks) and log the temperature during the timing run so a thermal excursion doesn’t masquerade as a slow model.
  • TensorRT engines are not portable. An engine built on one JetPack/TensorRT version and GPU will not load on another — rebuild from the ONNX on the target. Keep the ONNX as the source of truth, the .engine as a build artifact.

Background

Features. Each utterance/tone is a fixed-size log-mel patch \(x \in \mathbb{R}^{M \times T}\) (\(M\) mel bands × \(T\) frames), computed exactly as in Lab 8.1. A small CNN (a few conv+pool layers, global pool, one dense layer) maps it to \(C\) class logits \(z \in \mathbb{R}^C\).

Softmax and cross-entropy. The logits become class probabilities via softmax,

\[ \hat{p}_c \;=\; \frac{e^{z_c}}{\sum_{c'=1}^{C} e^{z_{c'}}}, \]

and training minimizes the cross-entropy between the one-hot label \(y\) and \(\hat p\), averaged over the batch:

\[ \mathcal{L} \;=\; -\sum_{c=1}^{C} y_c \log \hat{p}_c \;=\; -\log \hat{p}_{c^\star}, \]

where \(c^\star\) is the true class. This is the negative log-likelihood of a categorical model — the Course 1 Weeks 5–6 probability view — and it is minimized by (stochastic) gradient descent, the Weeks 7–8 optimization view. The gradient at the logits is beautifully simple,

\[ \frac{\partial \mathcal{L}}{\partial z_c} \;=\; \hat{p}_c - y_c, \]

which is why softmax+cross-entropy is the default classification head.

Quantization. int8 inference maps a float tensor \(r\) to an 8-bit integer \(q\) via a scale \(s\) and zero-point \(z_0\): \(q = \operatorname{round}(r/s) + z_0\). Post-training quantization calibrates \(s\) per tensor from a small representative set; the model then runs in integer arithmetic, ~4× smaller and much faster on hardware with int8 units (the Orin’s tensor cores, the Pi’s NEON). The cost is a small accuracy drop you must measure, not assume.

Why compare to Goertzel. Goertzel (Lab 6.5) computes a single DFT bin’s energy in \(O(N)\) with a two-tap recurrence — for a known tone frequency it is optimal, deterministic, and needs no training data. The CNN earns its complexity only when the classes are not cleanly separable by a few fixed frequencies: real spoken words, noisy conditions, speaker variation. Frame the lab as: Goertzel for DTMF, CNN for keywords, and measure both on the same tones to see the crossover.

Procedure

Part A — Dataset.

  1. For DTMF: synthesize labeled tone pairs from the MCP4725/STM32 across the standard keypad frequencies; capture them through the real mic so the model sees the real channel. For keywords: record a few words (e.g. “yes/no/stop/go”) in your own voice, many repetitions, varied distance/room, plus a “background/unknown” class. Hold out a test split recorded on a different day.
  2. Featurize every clip to a fixed \(M \times T\) log-mel patch (Lab 8.1 pipeline). Standardize (per-band mean/var from the training set).

Part B — Train on the host.

  1. Define a small CNN in PyTorch (aim for well under ~100k params — it has to fit an edge budget). Train with Adam, cross-entropy loss, early stopping on the held-out split. Log train/val loss and confusion matrix.
# illustrative model — owner writes and trains the real one
class KWS(nn.Module):
    def __init__(self, C):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1))
        self.fc = nn.Linear(32, C)
    def forward(self, x):            # x: (B,1,M,T)
        return self.fc(self.net(x).flatten(1))

Part C — Export and quantize.

  1. Export to ONNX (torch.onnx.export, fixed input shape). This ONNX is the single source of truth for both boards.
  2. Jetson: build an int8 TensorRT engine with a calibration set of ~a few hundred representative log-mel patches (trtexec --onnx=kws.onnx --int8 --calib=... or the Python calibrator API). Also build an fp16 engine for comparison.
  3. Pi 5: run the ONNX directly in onnxruntime, and/or convert to tflite with int8 post-training quantization (TF converter with a representative dataset). NEON/XNNPACK gives the int8 speedup on the A76.

Part D — Measure on-device.

  1. On each board, run the test split through the board’s own audio path and record: top-1 accuracy, confusion matrix, and per-inference latency (mean and p95 over ≥1000 runs, after warm-up and with clocks fixed). Separate feature latency (STFT+mel) from model latency — on the Pi the FFT can dominate.
  2. Run Goertzel (Lab 6.5) on the DTMF test set on the STM32 (or a Python port on the Pi) and record its accuracy and latency for the head-to-head.

Deliverable & expected results

  • Trained model + ONNX + a TensorRT int8 engine + a tflite/onnx int8 model.
  • An on-device results table (per board, per precision) and a confusion matrix.
  • A one-paragraph verdict on Goertzel-vs-CNN for the tone task.
Quantity Predicted / target Measured
Host fp32 top-1 accuracy (from training)
Jetson int8 top-1 accuracy ≈ host − small drop
Pi 5 int8 (tflite) top-1 accuracy ≈ host − small drop
Model latency, Jetson int8 (p95) (µs)
Model latency, Pi 5 int8 (p95) (µs)
Feature (STFT+mel) latency, Pi 5 (µs)
Goertzel latency per tone, STM32 \(O(N)\), a few µs
Goertzel DTMF accuracy ≈ 100% clean

For real-time keyword spotting the whole feature+model pipeline must fit inside one hop’s worth of time (16 ms/frame at \(f_s=16\) kHz, \(H=256\) from Lab 8.1) if you infer every frame, or the utterance window if you infer per utterance.

Analysis & reconciliation

Explain the accuracy gap between host fp32 and on-device int8 as quantization error plus mic-channel mismatch; if the drop is large (> a few points), your calibration set was probably unrepresentative or a layer is quantization-sensitive (try keeping the first/last layer in fp16). Compare latencies: expect the Jetson int8 engine to have the lowest model latency but a real fixed cost in host↔︎device copy and TensorRT launch overhead, and the Pi 5 to be competitive at these tiny model sizes because there’s no PCIe copy and the feature extraction is the real bottleneck. Then land the classical comparison honestly: for clean DTMF, Goertzel wins on every axis — lower latency, zero training data, deterministic, trivially verifiable — and the CNN only justifies itself on the genuinely spoken-keyword task where fixed-frequency detection fails. Carry the per-inference latency and accuracy numbers into Lab 8.5.

Going further

  • Add noise augmentation at training time (mix in the noise you’ll suppress in Lab 8.3) and re-measure on-device robustness.
  • Try quantization-aware training (QAT) and see how much of the int8 accuracy drop it recovers versus post-training quantization.
  • Measure the accuracy/latency Pareto as you shrink the model (channels, layers) — find the smallest net that still beats Goertzel on the keyword task.
  • Stream continuous audio and add a simple posterior smoothing / debounce so a single noisy frame doesn’t trigger a false keyword — the same debounce logic you’d write in firmware.