[Paper Review] nnAudio: An on-the-fly GPU Audio to Spectrogram Conversion Toolbox Using 1D Convolution Neural Networks

Kin Wai Cheuk, Hans Anderson+2 more|arXiv (Cornell University)|Dec 27, 2019

Music and Audio Processing66 references4 citations

TL;DR

nnAudio is a PyTorch-based, GPU-accelerated audio processing toolbox that uses 1D convolutional neural networks to perform on-the-fly, differentiable spectrogram conversion from raw waveforms to STFT, Mel spectrograms, and CQTs. It reduces spectrogram extraction time from seconds (librosa) to milliseconds, enables end-to-end training of transformation layers, and achieves up to 1000x speedup on GPU over CPU-based methods.

ABSTRACT

Converting time domain waveforms to frequency domain spectrograms is typically considered to be a prepossessing step done before model training. This approach, however, has several drawbacks. First, it takes a lot of hard disk space to store different frequency domain representations. This is especially true during the model development and tuning process, when exploring various types of spectrograms for optimal performance. Second, if another dataset is used, one must process all the audio clips again before the network can be retrained. In this paper, we integrate the time domain to frequency domain conversion as part of the model structure, and propose a neural network based toolbox, nnAudio, which leverages 1D convolutional neural networks to perform time domain to frequency domain conversion during feed-forward. It allows on-the-fly spectrogram generation without the need to store any spectrograms on the disk. This approach also allows back-propagation on the waveforms-to-spectrograms transformation layer, which implies that this transformation process can be made trainable, and hence further optimized by gradient descent. nnAudio reduces the waveforms-to-spectrograms conversion time for 1,770 waveforms (from the MAPS dataset) from $10.64$ seconds with librosa to only $0.001$ seconds for Short-Time Fourier Transform (STFT), $18.3$ seconds to $0.015$ seconds for Mel spectrogram, $103.4$ seconds to $0.258$ for constant-Q transform (CQT), when using GPU on our DGX work station with CPU: Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz Tesla v100 32Gb GPUs. (Only 1 GPU is being used for all the experiments.) We also further optimize the existing CQT algorithm, so that the CQT spectrogram can be obtained without aliasing in a much faster computation time (from $0.258$ seconds to only $0.001$ seconds).

Motivation & Objective

• To address the inefficiency and high storage cost of pre-computing and storing multiple spectrogram types during neural network hyperparameter search.
• To enable on-the-fly spectrogram extraction directly on GPU, eliminating disk I/O bottlenecks and CPU-GPU data transfer overhead.
• To make the spectrogram transformation layer trainable within the neural network, allowing optimization of Fourier kernels, Mel filter banks, and CQT kernels for task-specific performance.
• To provide a fast, scalable, and user-friendly GPU-based audio processing library that supports STFT, Mel spectrograms, and CQT with end-to-end differentiability.
• To outperform existing CPU and GPU-based libraries (e.g., librosa, Kapre, torchaudio) in processing speed while supporting novel differentiable CQT implementations.

Proposed method

• Leverages 1D convolutional neural networks in PyTorch to implement time-domain to frequency-domain transformations, treating STFT, Mel, and CQT as learnable layers.
• Uses CUDA-enabled 1D convolutions to perform spectrogram computation directly on GPU, avoiding CPU-GPU data transfer and enabling real-time processing.
• Represents Fourier kernels, Mel filter banks, and CQT kernels as learnable parameters that can be updated during backpropagation.
• Implements CQT using time-domain kernels, which are shown to be faster than traditional frequency-domain kernel methods.
• Supports multiple spectrogram types (STFT, Mel, CQT) with configurable parameters, all computed in a single, unified, differentiable framework.
• Optimizes memory usage by avoiding disk storage of intermediate spectrograms, enabling on-the-fly generation during training.

Experimental results

Research questions

• RQ1 Can a deep learning-based, GPU-accelerated framework achieve significantly faster spectrogram extraction than CPU-based libraries like librosa?
• RQ2 Can the spectrogram transformation layer be made trainable within a neural network, leading to improved downstream model performance?
• RQ3 Can differentiable CQT computation be efficiently implemented using 1D convolutions on GPU, and does it outperform existing frequency-domain kernel methods?
• RQ4 Does on-the-fly spectrogram generation with a trainable transformation layer reduce training time and improve model accuracy compared to precomputed, non-trainable spectrograms?
• RQ5 How does the proposed framework scale across different hardware platforms (e.g., consumer GPUs vs. high-end DGX systems) in terms of speed and memory efficiency?

Key findings

• nnAudio reduces STFT extraction time from 10.6 seconds (librosa on CPU) to 0.001 seconds on GPU, achieving a 10,000x speedup.
• Mel spectrogram computation is accelerated from 18.3 seconds (librosa) to 0.015 seconds with nnAudio, a 1,200x improvement.
• CQT processing time drops from 103.4 seconds (librosa) to 0.258 seconds with nnAudio, a 400x speedup.
• On the MusicNet dataset, nnAudio reduces total spectrogram extraction time from 983 minutes (librosa) to 99 minutes, a 90% reduction.
• The use of time-domain CQT kernels in nnAudio enables faster computation than traditional frequency-domain kernel methods, reducing CQT time from 0.258s to 0.001s for 1,770 waveforms.
• Training the spectrogram transformation layer (e.g., Fourier kernels, Mel filters) leads to a lower MSE loss in a frequency prediction task compared to non-trainable kernels, demonstrating the benefit of end-to-end optimization.