Lab 8.1 — GPU Real-Time Spectrogram (Jetson)

← Course 2 syllabus · Bonus Module 8 · Prev: « Lab 7.3 · Next: Lab 8.2 »

Goal

Take the same audio-band signals you have been generating and measuring all course and compute a real-time short-time Fourier transform (STFT) and log-mel spectrogram on the Jetson’s GPU. The skill built is the front end of essentially every modern audio-ML system: framing, windowing, batched FFT, and a mel filterbank, done at streaming rates with no dropped frames. Doing it on the GPU (CuPy or torchaudio) and timing it against the same pipeline on the Pi 5 CPU (NumPy) teaches you where an accelerator actually pays off — and where the PCIe/USB and host-to-device copy costs eat the win. This is the bridge from the classical STM32 FFT of Lab 6.3 to learned DSP: the log-mel spectrogram is the input feature for Labs 8.2–8.4.

Recommended reading

• Lyons Ch. 3 — the DFT: leakage, windowing, and the picket-fence effect. This is the intuition for why we window each STFT frame. → ../books/lyons-dsp/index.qmd
• Hayes Ch. 8 — spectrum estimation: the periodogram, Welch’s method, bias/variance of spectral estimates. The STFT is a time-localized periodogram. → ../books/hayes-statistical-dsp/index.qmd
• Course 1 Week 10 — the Fourier transform of tempered distributions and the sampling theorem; the rigorous reason a discrete spectrum represents a continuous one. → ../course1/index.qmd#week-10
• Optional: the torchaudio.transforms.MelSpectrogram and CuPy cupyx.scipy.fft docs for the exact API surface you will call.

Equipment & parts

• Jetson Orin Nano with JetPack (CUDA/cuDNN, PyTorch, CuPy installed).
• Raspberry Pi 5 (for the CPU-vs-GPU comparison), Python with NumPy/SciPy.
• An audio input into each board — one of:
  • a USB microphone (simplest; class-compliant, appears as an ALSA capture device), or
  • an I²S MEMS mic on the 40-pin header, or
  • the ADS1115 reading an analog mic/line signal (low rate: ≤860 SPS, so only for sub-400 Hz tones — fine for the DTMF/tone work, not for speech).
• Your existing signal sources for a known input: the MCP4725 DAC (Lab 3.3) or STM32 DAC/PWM driving a tone into the mic or line input, so you can verify a bin lands where the math says.

Safety & don’t-break-it

• Line levels, not headphone-amp levels, into the ADC. If you feed an analog audio signal into the ADS1115, it must stay within GND−0.3 V to VDD+0.3 V. A mic preamp or line output can swing negative; bias it to mid-rail and/or clamp before the ADC input. Never feed a raw ±1 V line signal into a single-supply ADC pin.
• The Jetson and Pi are 3.3 V-logic boards. I²S and GPIO pins are not 5 V tolerant. Level-shift (Lab 3.5) anything coming from a 5 V part.
• Common ground between the audio source, the ADC/mic, and the board — a floating mic ground shows up as hum and a huge DC bin.
• Thermals. Sustained GPU FFT work heats the Orin Nano. Run it with its heatsink/fan and in a case with airflow; a throttled GPU will quietly wreck your timing numbers.
• USB mics enumerate at 5 V bus power — fine, but don’t hot-plug during a capture run or you’ll drop the stream.

Background

Framing and the STFT. Split the sampled signal \(x[n]\) into overlapping frames of length \(N\) (the window length), advanced by a hop \(H\) samples. Multiply each frame by a window \(w[n]\) (Hann is the default) and take its DFT. The STFT is

\[ X[m, k] \;=\; \sum_{n=0}^{N-1} w[n]\,x[n + mH]\; e^{-j\,2\pi k n / N}, \qquad k = 0,\dots,N-1, \]

where \(m\) indexes the frame (time) and \(k\) indexes the frequency bin. Bin \(k\) sits at frequency \(f_k = k f_s / N\), and each frame covers \(N/f_s\) seconds. The frame rate is \(f_s / H\); the overlap is \(1 - H/N\) (75% overlap, \(H = N/4\), is a common default). The window trades main-lobe width (frequency resolution) against side-lobe leakage — the Lyons Ch. 3 material.

Power and log scaling. The spectrogram magnitude is \(|X[m,k]|^2\); we display it in dB, \(10\log_{10}(|X[m,k]|^2 + \epsilon)\), with a small \(\epsilon\) floor so silence doesn’t go to \(-\infty\).

The mel filterbank. Human pitch perception is roughly logarithmic, so audio-ML compresses the linear FFT bins into \(M\) mel bands. The mel scale is

\[ m(f) \;=\; 2595 \, \log_{10}\!\left(1 + \frac{f}{700}\right), \]

and its inverse maps \(M+2\) equally-spaced mel points back to Hz to place the triangular filters. Filter \(j\) has response \(H_j[k]\) (a triangle peaking at its center bin, zero at its neighbors), and the mel energy is

\[ S[m, j] \;=\; \sum_{k} H_j[k]\,|X[m,k]|^2 . \]

The log-mel spectrogram is \(\log(S[m,j] + \epsilon)\) — a small \(M \times (\text{frames})\) image, the standard learned-DSP feature. Everything here is linear in \(|X|^2\) except the final log, so the whole thing is a batched matrix multiply — exactly what a GPU is built for: stack all frames into an \(N \times (\text{frames})\) matrix, one batched FFT, then multiply by the fixed \(M \times (N/2{+}1)\) mel matrix.

Why the GPU may or may not win. The FFT of one 1024-point frame is tiny; the GPU wins only when you batch hundreds of frames per launch and keep data resident on the device. If you copy one frame at a time host→device→host, the PCIe/copy latency dominates and the Pi 5’s NumPy can be faster. That trade-off is the whole point of the timing comparison.

Procedure

Part A — Capture a known tone (both boards).

1. Drive a clean 440 Hz (or a DTMF pair) tone from the MCP4725/STM32 into the audio input. Pick a sample rate \(f_s\): 16 kHz for speech-style work (USB/I²S mic), or 8 kHz if you only care about tones.
2. Capture ~2 s into a NumPy array with sounddevice (sd.rec) or an ALSA read. Confirm the level: peak should be well below full scale, no clipping.

Part B — CPU reference STFT (Pi 5, NumPy).

3. Choose \(N = 1024\), \(H = 256\) (75% overlap), Hann window. Frame the signal (np.lib.stride_tricks.sliding_window_view then stride by \(H\)), apply the window, and np.fft.rfft along the frame axis.
4. Form the power spectrogram and the log-mel (build the mel matrix once with librosa.filters.mel or by hand from the formula above). Time the STFT+mel over the whole clip; divide by the number of frames to get µs/frame.

# illustrative — owner writes the real pipeline
frames = win[None, :] * sliding[:, ::H, :]        # (F, N) windowed frames
spec   = np.abs(np.fft.rfft(frames, axis=-1))**2  # (F, N/2+1) power
logmel = np.log(spec @ mel_fb.T + 1e-6)           # (F, M) log-mel

Part C — GPU STFT (Jetson, CuPy or torchaudio).

5. Repeat with the same \(N, H, M\). Two clean options:
   • CuPy: move the frame matrix to the device (cupy.asarray), cupyx.scipy.fft.rfft, matmul against a device-resident mel matrix, cupy.asnumpy only the final small log-mel back.
   • torchaudio: MelSpectrogram(sample_rate, n_fft=N, hop_length=H, n_mels=M).cuda() and push the waveform tensor to cuda.
6. Warm up first (run the transform once and cupy.cuda.Stream.null.synchronize() / torch.cuda.synchronize() before timing — the first launch pays JIT/cuFFT-plan cost). Then time a synchronized batched run and compute µs/frame.

Part D — Streaming, real-time.

7. Now do it live: a sounddevice input callback fills a ring buffer; a worker pulls \(N\)-sample frames every \(H\) samples and pushes them through the GPU transform in batches (e.g. accumulate 32 frames, transform once). Display the rolling log-mel with matplotlib (imshow with origin='lower', updated in place) or write it to disk.
8. Verify no overruns: the callback must never block on the GPU. Keep the FFT worker on its own thread and the device copies asynchronous.

Deliverable & expected results

• A rolling log-mel spectrogram of live audio on the Jetson, plus a saved PNG of a known tone showing the bin at the expected \(f_k\).
• A timing table: µs/frame for the STFT (and for STFT+mel) on Pi 5 (NumPy) vs Jetson (single-frame) vs Jetson (batched, device-resident).

For \(f_s = 16\text{ kHz}\), \(N = 1024\): bin spacing \(f_s/N = 15.625\text{ Hz}\), so a 440 Hz tone lands in bin \(k = 440 / 15.625 \approx 28\). Frame duration \(N/f_s = 64\text{ ms}\); at \(H = 256\) the frame rate is \(f_s/H = 62.5\) frames/s (16 ms/frame), so real time means the pipeline must sustain < 16 ms/frame end to end.

Quantity	Predicted	Measured
Bin spacing \(f_s/N\) (16 kHz, 1024)	15.625 Hz	…
Bin index of a 440 Hz tone	≈ 28	…
Frame rate at \(H=256\)	62.5 fps (16 ms/frame)	…
STFT+mel, Pi 5 NumPy	(µs/frame)	…
STFT+mel, Jetson single-frame	(µs/frame)	…
STFT+mel, Jetson batched (device-resident)	(µs/frame)	…

Analysis & reconciliation

Confirm the tone lands in the predicted bin — if it’s off by one or two bins, check your actual \(f_s\) (USB mics rarely run at exactly the requested rate) and whether you used rfft bin indexing consistently. Explain your timing result in terms of the copy vs compute split: expect the single-frame GPU path to lose or barely tie the Pi 5 CPU because the host↔︎device copy dominates a 1024-point FFT, and the batched, device-resident GPU path to win decisively as batch size grows (throughput up, per-frame overhead amortized). If the batched GPU path doesn’t pull ahead, you are probably still copying per frame or not warming up the cuFFT plan. Note the latency-vs-throughput tension: batching lowers µs/frame but raises the latency of any single frame — you carry this observation forward to the benchmark in Lab 8.5.

Going further

• Replace the mel filterbank with a learnable front end (a 1-D conv over the waveform, SincNet-style) and compare its learned filters to the fixed mel triangles — the theme of the rest of Module 8.
• Sweep \(N\) and \(H\) and plot the time-frequency resolution trade-off (the uncertainty principle from Course 1 Week 10) directly on the spectrogram.
• Compare torchaudio fp32 vs fp16 STFT timing on the Jetson — fp16 halves the memory traffic and is the first easy accelerator win.
• Feed the STM32 FFT output (Lab 6.3) over UART and overlay it on the GPU spectrogram of the same tone to sanity-check all three implementations against each other.