如何将.wav音频文件加载为tensorflow.data.Dataset格式数据集

前置依赖 #

仅需要安装TensorFlow 2.x版本即可，导入所需模块：

import tensorflow as tf
import os
from glob import glob # 也可使用glob更便捷地匹配wav文件

具体实现步骤 #

• 步骤1：获取所有wav文件的路径列表
  不需要提前加载音频内容，仅读取文件路径即可，内存开销可忽略：

  # 替换为你的音频文件存储目录
  WAV_DIR = "./your_wav_folder"
  # 匹配目录下所有wav文件路径
  wav_paths = glob(os.path.join(WAV_DIR, "*.wav"))
  

• 步骤2：编写单音频加载与预处理函数
  该函数会被Dataset流水线自动调用，仅在需要用到对应音频时才执行加载：

  def load_and_preprocess_wav(file_path):
      # 读取音频文件二进制内容
      wav_binary = tf.io.read_file(file_path)
      # 解码wav，返回值为(音频张量, 采样率)
      wav, sr = tf.audio.decode_wav(wav_binary, desired_channels=1) # desired_channels=1直接转单声道，不需要后续处理
      
      # 以下为可选预处理步骤，根据你的autoencoder需求调整
      # 1. 固定音频长度：超过则截断，不足则补0
      target_len = 16000 * 3 # 示例为3秒、16k采样率的音频长度
      wav = tf.squeeze(wav, axis=-1) # 去掉通道维度，形状从(seq_len, 1)变为(seq_len,)
      wav = wav[:target_len] if tf.shape(wav)[0] >= target_len else tf.pad(wav, [[0, target_len - tf.shape(wav)[0]]])
      # 2. 数值归一化到[0,1]或者[-1,1]，这里示例和decode_wav输出范围一致为[-1,1]
      # 3. 如需统一采样率，可使用tf.audio.resample(wav, input_rate=sr, output_rate=16000)
      
      # 因为autoencoder是自监督训练，输入和标签都是音频本身，所以返回(wav, wav)，如果不需要标签直接返回wav即可
      return wav, wav
  

• 步骤3：构建流式加载的Dataset流水线
  全程懒加载，不会一次性把所有音频加载到内存，完全不会出现内存溢出崩溃问题：

  batch_size = 8 # 根据你的显存/内存大小调整
  
  dataset = tf.data.Dataset.from_tensor_slices(wav_paths)
  # 并行加载预处理，速度更快
  dataset = dataset.map(load_and_preprocess_wav, num_parallel_calls=tf.data.AUTOTUNE)
  # 训练集可选操作：打乱顺序
  dataset = dataset.shuffle(buffer_size=len(wav_paths)//10) # buffer_size不要设太大，避免占用过多内存
  # 分批
  dataset = dataset.batch(batch_size)
  # 预加载下一批次数据，训练时无IO等待
  dataset = dataset.prefetch(tf.data.AUTOTUNE)
  
  # 测试数据集是否正常
  for batch_x, batch_y in dataset.take(1):
      print(batch_x.shape, batch_y.shape)
  

额外优化建议 #

• 如果你的音频数量超过1万，可在map步骤后添加cache("./wav_cache")，将预处理后的音频缓存到磁盘，第二次及之后训练不需要重复预处理，速度会快很多，且不会占用内存
• 如果音频采样率较高、单文件较大，可以适当调小batch_size避免内存溢出
• 如果是测试集不需要打乱，可以去掉shuffle步骤

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