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 Lesson 49 — the Fourier transform of tempered distributions and the sampling theorem; the rigorous reason a discrete spectrum represents a continuous one. → ../course1/index.qmd#lesson-49
- Optional: the
torchaudio.transforms.MelSpectrogramand CuPycupyx.scipy.fftdocs 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.
Wiring & bench setup
No breadboard for the recommended path: a USB microphone plugs straight into a USB-A port on each board and enumerates as an ALSA capture device (arecord -l to find its card number). Plug it into the board directly, not through an unpowered hub, and don’t re-plug it mid-capture. The known-tone source (MCP4725/STM32 DAC) plays acoustically into the mic through a powered speaker, or couples into the line/ADC input (mid-rail-biased — see Safety).
If you take the I²S MEMS mic option (INMP441/SPH0645-class breakout), it wires to the 40-pin header — same header positions on the Pi 5 and the Jetson Orin Nano; 3.3 V logic only:
| Mic breakout pin | Purpose | 40-pin header |
|---|---|---|
| VDD | 3.3 V | pin 1 (3V3) |
| GND | ground | pin 6 (GND) |
| SCK / BCLK | bit clock | pin 12 |
| WS / LRCL | word select | pin 35 |
| SD / DOUT | data → board | pin 38 |
On the Pi, enable the I²S overlay in /boot/firmware/config.txt per your breakout’s guide before it will capture. The ADS1115 option wires exactly as in Lab 3.4, on the board’s I²C header pins (SDA = pin 3, SCL = pin 5).
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.
Project & environment setup
Both boards run their scripts out of the lab’s edge/ folder — sync the repo (or just labs/) to each board so results land in labs/lab-8-1/edge/ and come back with the repo.
Pi 5 (Raspberry Pi OS Bookworm — pip only works inside a venv):
sudo apt install libportaudio2 # PortAudio backend for sounddevice
python3 -m venv ~/edge-venv && source ~/edge-venv/bin/activate
pip install numpy scipy sounddevice matplotlib librosa # librosa only for filters.mel — or build the mel matrix by hand
mkdir -p labs/lab-8-1/edge labs/lab-8-1/capturesThe SciPy piece this lab implies is scipy.signal (scipy.signal.stft/spectrogram as a cross-check on your stride-tricks STFT, scipy.signal.get_window for the Hann window) — owner writes the code.
Jetson Orin Nano — JetPack already ships CUDA, cuDNN, and TensorRT; you add only the Python compute stack:
sudo apt install libportaudio2
python3 -m venv ~/edge-venv && source ~/edge-venv/bin/activate
pip install numpy sounddevice matplotlib
pip install cupy-cuda12x # CuPy path — matches JetPack's CUDA 12 (or NVIDIA's Jetson-specific wheel per their docs)On the Jetson, cupyx.scipy.signal (with cupyx.scipy.fft) is CuPy’s drop-in GPU mirror of the scipy.signal API — same calls, swap the import, and the arrays live on-device as cupy arrays.
For the torchaudio path instead, install NVIDIA’s Jetson PyTorch wheel (the plain PyPI torch wheel has no CUDA on aarch64 — see NVIDIA’s “PyTorch for Jetson” page), then a matching torchaudio.
Scripts (you write them) live under labs/lab-8-1/edge/: stft_cpu.py (Part B NumPy reference), stft_gpu.py (Part C CuPy/torchaudio + timing), stream_logmel.py (Part D ring-buffer streamer).
Where results go:
| Artifact | Path |
|---|---|
| Bench note (timing table filled in) | labs/lab-8-1/notes.md |
| Captured known-tone clip | labs/lab-8-1/captures/tone-440.wav |
| Log-mel PNG of the known tone | labs/lab-8-1/edge/logmel-tone.png |
| µs/frame timings (both boards, all three paths) | labs/lab-8-1/edge/timing.csv |
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).
- Drive a clean 440 Hz (or a DTMF pair) tone from the MCP4725/STM32 into the audio input (per Wiring & bench setup). 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.
- 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).
- Choose \(N = 1024\), \(H = 256\) (75% overlap), Hann window. Frame the signal (
np.lib.stride_tricks.sliding_window_viewthen stride by \(H\)), apply the window, andnp.fft.rfftalong the frame axis. - Form the power spectrogram and the log-mel (build the mel matrix once with
librosa.filters.melor 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-melPart C — GPU STFT (Jetson, CuPy or torchaudio).
- 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.asnumpyonly 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 tocuda.
- CuPy: move the frame matrix to the device (
- 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.
- Now do it live: a
sounddeviceinput 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 (imshowwithorigin='lower', updated in place) or write it to disk. - 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 time–frequency trade-off of Course 1 Lesson 36) directly on the spectrogram.
- Compare
torchaudiofp32 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.