Lab 8.1 — GPU Real-Time Spectrogram (Jetson)
← Course 3 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.
- Hayes Ch. 8 — spectrum estimation: the periodogram, Welch’s method, bias/variance of spectral estimates. The STFT is a time-localized periodogram.
- Course 1 Lesson 35 — the Fourier transform of tempered distributions and the sampling theorem; the rigorous reason a discrete spectrum represents a continuous one. → Course 1 Lesson 35
- 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:
- the USB sound card’s onboard microphone (on hand — see the bench inventory; class-compliant, appears as an ALSA capture device: find its card index with
arecord -land pass it tosounddevice/arecord), 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).
- the USB sound card’s onboard microphone (on hand — see the bench inventory; class-compliant, appears as an ALSA capture device: find its card index with
- 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
Acoustic loopback. Play the known tone through the sound card’s own speakers and record it with its microphone in the same process (
sounddevice.playrec); the spectrogram now contains the speaker, the room, and the mic — compare the recorded bin amplitude against the electrical one and you have a first frequency-response measurement of the whole audio chain.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 31) 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.