如何解决Librosa批量处理音频时的内存泄漏问题？

问题背景 #

需处理约700个.wav音频文件，转换为频谱图并保存为.jpg格式。每次迭代后手动删除变量，但内存占用仍持续上升，处理300个文件后内存达30GB，导致系统崩溃。原代码如下：

FRAME_SIZE=2048
HOP_SIZE=512
DURATION=30
def save_and_plot_spectrogram(Y, out_dir, sr, hop_length=HOP_SIZE, y_axis="log", cmap='viridis'):
    plt.figure(figsize=(25, 10))
    librosa.display.specshow(Y,
                             sr=sr,
                             hop_length=hop_length,
                             x_axis="time",
                             y_axis=y_axis,
                             cmap=cmap)  # Set the colormap here
    plt.axis('off')
    plt.savefig(out_dir, bbox_inches='tight', pad_inches=0)
    plt.close()  # Close the plot to free memory
    plt.clf()    # Clear the current figure
    plt.cla()    # Clear the current axes


def process_signal(path,out_dir):
    x,sr = librosa.load(path,duration=DURATION)
    x_stft=librosa.stft(x,n_fft=FRAME_SIZE,hop_length=HOP_SIZE)
    x_mag=np.abs(x_stft)**2
    x_log=librosa.power_to_db(x_mag)

    save_and_plot_spectrogram(x_log,out_dir,sr)
    
    del x
    del x_stft
    del x_mag
    del x_log
    
def extract_sprectrogram_from_audio(source_folder,spectogram_folder):
    os.makedirs(spectogram_folder, exist_ok=True)

    for filename in os.listdir(source_folder):
        audio_path=os.path.join(source_folder,filename)

        audio_name=os.path.splitext(filename)[0]
        audio_spectogram_path=audio_name+'.jpg'
        output_dir =os.path.join(spectogram_folder,audio_spectogram_path)
        process_signal(audio_path,output_dir)


source_folder='/kaggle/input/audio-dataset/wav'
spectogram_folder='/kaggle/working/spectogram'
extract_sprectrogram_from_audio(source_folder,spectogram_folder)

内存泄漏原因 #

• Matplotlib资源未彻底释放：save_and_plot_spectrogram中plt.close()后调用plt.clf()、plt.cla()属于冗余操作，且未显式引用创建的figure对象，导致部分绘图资源残留。
• 垃圾回收未及时触发：手动del变量后，Python垃圾回收器不会立即回收大数组内存，循环中积累的未释放内存持续占用空间。
• 循环临时对象累积：os.listdir生成的文件路径字符串等临时对象在全局作用域残留，未被及时清理。

修复方案 #

1. 彻底清理Matplotlib绘图资源 #

显式创建并引用figure对象，关闭后销毁该对象，避免资源残留：

def save_and_plot_spectrogram(Y, out_dir, sr, hop_length=HOP_SIZE, y_axis="log", cmap='viridis'):
    # 显式创建figure并持有引用
    fig = plt.figure(figsize=(25, 10))
    librosa.display.specshow(Y,
                             sr=sr,
                             hop_length=hop_length,
                             x_axis="time",
                             y_axis=y_axis,
                             cmap=cmap)
    plt.axis('off')
    plt.savefig(out_dir, bbox_inches='tight', pad_inches=0)
    # 关闭指定figure并销毁引用
    plt.close(fig)
    del fig

2. 手动触发垃圾回收 #

在处理函数末尾强制触发垃圾回收，立即释放未使用的内存：

import gc

def process_signal(path,out_dir):
    x,sr = librosa.load(path, duration=DURATION, mono=True, dtype=np.float32)
    x_stft=librosa.stft(x,n_fft=FRAME_SIZE,hop_length=HOP_SIZE)
    x_mag=np.abs(x_stft)**2
    x_log=librosa.power_to_db(x_mag)

    save_and_plot_spectrogram(x_log,out_dir,sr)
    
    del x, x_stft, x_mag, x_log
    # 强制回收内存
    gc.collect()

3. 优化循环与变量作用域 #

使用pathlib遍历文件，减少临时字符串对象，同时在循环末尾也触发垃圾回收：

from pathlib import Path

def extract_sprectrogram_from_audio(source_folder,spectogram_folder):
    os.makedirs(spectogram_folder, exist_ok=True)
    source_path = Path(source_folder)
    # 仅遍历wav文件，避免无效处理
    for audio_path in source_path.glob("*.wav"):
        audio_name = audio_path.stem
        output_dir = Path(spectogram_folder) / f"{audio_name}.jpg"
        process_signal(str(audio_path), str(output_dir))
        # 循环内回收内存
        gc.collect()

4. 降低音频加载内存占用 #

加载音频时指定dtype=np.float32，将默认的float64改为float32，直接减少一半内存占用：

x,sr = librosa.load(path, duration=DURATION, mono=True, dtype=np.float32)

内容的提问来源于stack exchange，提问作者Vinayak Bansal