Whisper-Diarization项目处理长音频时显存不足问题的解决方案

问题背景

在使用Whisper-Diarization项目进行语音转写和说话人分离时，处理长时间音频文件（特别是超过2小时）经常会遇到CUDA显存不足的问题。这个问题尤其容易出现在多说话人场景中，当音频包含15个以上说话人时，显存需求会显著增加。

问题分析

经过技术分析，发现该问题主要由以下几个因素导致：

1. 说话人数量过多：默认配置下，MSDD模型最多支持8个说话人，当实际说话人数量超过这个限制时，显存消耗会急剧增加。

2. 模型处理机制：在说话人分离阶段，模型需要同时处理多个尺度的音频特征和说话人嵌入向量，这会占用大量显存资源。

3. 批量处理设置：默认的批量处理大小(infer_batch_size)可能不适合所有硬件配置，特别是对于长时间多说话人音频。

解决方案

1. 调整关键配置参数

在项目配置文件nemo_msdd_configs/diar_infer_telephonic.yaml中，可以调整以下参数来优化显存使用：

diarizer:
  msdd_model:
    parameters:
      infer_batch_size: 16  # 降低批量处理大小
      chunk_cluster_count: 10  # 减少每个音频块的聚类数量
      embeddings_per_chunk: 1000  # 减少每个块处理的嵌入向量数量
      max_num_speakers: 8  # 设置最大说话人数量

2. 参数调整建议

• infer_batch_size：这是最关键参数，建议从16开始尝试，根据显存情况逐步调整
• embeddings_per_chunk：可以从默认的10000降低到1000或更低
• chunk_cluster_count：最小可设置为max_num_speakers+2

3. 其他优化策略

如果参数调整后仍遇到显存问题，可以考虑：

1. 音频分割：将长音频分割成30-60分钟的片段分别处理
2. 硬件选择：使用具有更大显存的GPU（建议至少12GB显存）
3. 模型选择：尝试使用较小规模的Whisper模型

实施效果

在实际测试中，对一个2小时包含15个说话人的音频文件：

• 默认配置下显存消耗达到17GB，导致OOM错误
• 将infer_batch_size调整为16后，峰值显存使用降至15GB
• 处理过程顺利完成，识别准确率无明显下降

最佳实践建议

1. 对于多说话人场景，建议先估计说话人数量并相应调整max_num_speakers
2. 处理长音频时，优先调整infer_batch_size参数
3. 监控GPU显存使用情况，找到最适合自己硬件的参数组合
4. 考虑使用脚本自动化音频分割和结果合并流程

通过合理配置这些参数，可以在不显著影响识别质量的前提下，有效解决Whisper-Diarization项目处理长音频时的显存不足问题。