使用Speechbrain增强音频时CUDA内存不足的解决方法咨询

我尝试用Speechbrain对一段3分16秒的音频进行增强，运行Sepformer模型的代码时触发了CUDA内存不足错误：

OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 39.59 GiB total capacity; 33.60 GiB already allocated; 3.19 MiB free; 38.06 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

使用的代码如下：

from speechbrain.pretrained import SepformerSeparation as separator
import torchaudio

model = separator.from_hparams(source="speechbrain/sepformer-wham-enhancement",         
    savedir='pretrained_models/sepformer-wham-enhancement', run_opts={"device":"cuda"})

est_sources = model.separate_file(path=audio_file) 

torchaudio.save("enhanced_wham.wav", est_sources[:, :, 0].detach().cpu(), 8000)

请问推荐的解决方法是什么？是否应该将音频文件分割成片段处理？

核心解决方法 #

分割音频片段处理是非常推荐的方案，也是长音频处理中规避OOM的标准做法。除此之外，还有以下几种可行方案：

1. 分割音频片段处理

将长音频切割成固定时长的片段（比如10-30秒，具体时长可根据你的GPU显存调整），逐段处理后再拼接成完整音频。示例代码如下：

from speechbrain.pretrained import SepformerSeparation as separator
import torchaudio
import torch

# 加载模型
model = separator.from_hparams(source="speechbrain/sepformer-wham-enhancement",         
    savedir='pretrained_models/sepformer-wham-enhancement', run_opts={"device":"cuda"})

# 加载音频并转换采样率
waveform, sr = torchaudio.load(audio_file)
if sr != 8000:
    waveform = torchaudio.transforms.Resample(orig_freq=sr, new_freq=8000)(waveform)

# 定义片段参数
segment_duration = 20  # 单位：秒
segment_samples = segment_duration * 8000
total_samples = waveform.size(1)
enhanced_segments = []

# 逐段处理
for i in range(0, total_samples, segment_samples):
    end = min(i + segment_samples, total_samples)
    segment = waveform[:, i:end].to(model.device)
    est_source = model.separate_batch(segment)[:, :, 0].detach().cpu()
    enhanced_segments.append(est_source)

# 拼接并保存结果
enhanced_waveform = torch.cat(enhanced_segments, dim=1)
torchaudio.save("enhanced_wham.wav", enhanced_waveform, 8000)

2. 优化PyTorch内存配置

设置环境变量减少内存碎片，可在运行脚本前执行：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

也可以在Python代码开头添加：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

3. 切换到CPU推理

如果GPU显存不足，改用CPU处理（速度较慢但无内存限制），修改模型初始化参数：

model = separator.from_hparams(source="speechbrain/sepformer-wham-enhancement",         
    savedir='pretrained_models/sepformer-wham-enhancement', run_opts={"device":"cpu"})

4. 释放GPU冗余内存

用nvidia-smi命令检查GPU占用情况，关闭无关进程后重新运行脚本。

5. 半精度推理

使用float16半精度推理降低内存占用，修改模型初始化代码：

import torch
model = separator.from_hparams(source="speechbrain/sepformer-wham-enhancement",         
    savedir='pretrained_models/sepformer-wham-enhancement', 
    run_opts={"device":"cuda", "dtype":torch.float16})

注意：部分模型可能对低精度推理兼容性不佳，需验证输出效果。

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