[论文解读] 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 Processing参考文献 66被引用 4

一句话总结

nnAudio 是一个基于 PyTorch 的 GPU 加速音频处理工具箱，利用一维卷积神经网络实现从原始波形到 STFT、梅尔频谱图和 CQT 的实时、可微分的频谱图转换。它将频谱图提取时间从秒级（librosa）缩短至毫秒级，支持转换层的端到端训练，并在 GPU 上相比 CPU 方法实现了高达 1000 倍的加速。

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).

研究动机与目标

• 为解决在神经网络超参数搜索过程中预计算并存储多种频谱图类型所导致的低效与高存储成本问题。
• 实现在 GPU 上直接进行实时频谱图提取，消除磁盘 I/O 瓶颈和 CPU-GPU 数据传输开销。
• 使频谱图变换层在神经网络内可训练，从而对傅里叶核、梅尔滤波器组和 CQT 核进行任务特定的优化。
• 提供一个快速、可扩展且用户友好的基于 GPU 的音频处理库，支持 STFT、梅尔频谱图和 CQT，并具备端到端可微性。
• 在处理速度上超越现有 CPU 和 GPU 库（如 librosa、Kapre、torchaudio），同时支持新型可微分 CQT 实现。

提出的方法

• 利用 PyTorch 中的一维卷积神经网络实现时域到频域的转换，将 STFT、梅尔频谱图和 CQT 视为可学习层。
• 使用 CUDA 加速的一维卷积在 GPU 上直接执行频谱图计算，避免 CPU-GPU 数据传输，实现实时处理。
• 将傅里叶核、梅尔滤波器组和 CQT 核表示为可学习参数，可在反向传播过程中更新。
• 采用时域核实现 CQT，其速度优于传统的频域核方法。
• 支持多种频谱图类型（STFT、梅尔频谱图、CQT），参数可配置，所有计算均在一个统一、可微分的框架中完成。
• 通过避免中间频谱图的磁盘存储，优化内存使用，实现在训练过程中实时生成频谱图。

实验结果

研究问题

• RQ1 基于深度学习的 GPU 加速框架能否显著提升频谱图提取速度，超越 librosa 等 CPU 库？
• RQ2 能否在神经网络内实现频谱图变换层的可训练性，从而提升下游模型性能？
• RQ3 能否使用一维卷积在 GPU 上高效实现可微分的 CQT 计算？其性能是否优于现有频域核方法？
• RQ4 与预计算且不可训练的频谱图相比，采用可训练变换层的实时频谱图生成是否能减少训练时间并提高模型准确率？
• RQ5 在不同硬件平台（如消费级 GPU 与高端 DGX 系统）上，该框架在速度和内存效率方面的可扩展性如何？

主要发现

• nnAudio 将 STFT 提取时间从 CPU 上的 10.6 秒（librosa）缩短至 0.001 秒，实现 10,000 倍加速。
• 梅尔频谱图计算时间从 librosa 的 18.3 秒缩短至 nnAudio 的 0.015 秒，提升 1,200 倍。
• CQT 处理时间从 librosa 的 103.4 秒降至 nnAudio 的 0.258 秒，加速 400 倍。
• 在 MusicNet 数据集上，nnAudio 将总频谱图提取时间从 983 分钟（librosa）减少至 99 分钟，降低 90%。
• nnAudio 中采用的时域 CQT 核使计算速度优于传统频域核方法，将 1,770 个波形的 CQT 时间从 0.258 秒降至 0.001 秒。
• 训练频谱图变换层（如傅里叶核、梅尔滤波器）在频率预测任务中相比不可训练核，可实现更低的 MSE 损失，证明了端到端优化的优势。