使用audio_dataset_from_directory训练Keras模型时内核无响应求助

以下是针对训练卡顿、内核崩溃问题的具体解决思路：

1. 解决变长序列导致的内存/计算瓶颈 #

你使用ragged=True加载变长音频序列，但SimpleRNN处理批量变长序列时，易出现内存占用飙升或计算效率极低的情况。建议先将所有音频统一为固定长度：

import tensorflow as tf

def process_audio(audio, label):
    # 设定固定长度（示例：160000个采样点，对应10秒16kHz音频，根据你的数据调整）
    target_len = 160000
    audio_len = tf.shape(audio)[0]
    
    # 裁剪过长音频，填充过短音频
    if audio_len > target_len:
        audio = audio[:target_len]
    else:
        padding = target_len - audio_len
        audio = tf.pad(audio, [[0, padding]])
    
    # 保持形状为 (timesteps, 1)
    audio = tf.expand_dims(audio, axis=-1)
    return audio, label

# 加载数据集后添加处理步骤
data = data.map(process_audio, num_parallel_calls=tf.data.AUTOTUNE)

2. 减小批量大小控制内存占用 #

audio_dataset_from_directory默认batch_size为32，即使总音频文件只有320MB，转成张量后加上RNN的中间计算，内存消耗会远超文件大小。尝试将batch_size调小：

data = keras.utils.audio_dataset_from_directory(
    ".\\files\\songs", 
    labels=data["UserLikes"].to_list(), 
    label_mode="int", 
    ragged=True, 
    shuffle=True,
    batch_size=8  # 从8开始测试，可根据内存情况进一步调整
)

3. 预预处理音频并离线保存 #

实时加载和预处理音频会增加训练时的IO和计算开销，建议先将音频预处理后保存为numpy数组，后续直接加载训练：

import numpy as np

# 创建预处理数据保存目录
preprocessed_dir = ".\\preprocessed_songs"
os.makedirs(preprocessed_dir, exist_ok=True)

# 遍历音频文件完成预处理并保存
metadata = pd.read_csv(".\\files\\metadata.csv", index_col="title")
for root, _, files in os.walk(".\\files\\songs"):
    for file in files:
        if file.endswith((".wav", ".mp3")):
            # 加载音频并统一采样率（示例转成16kHz）
            audio, sr = keras.utils.load_audio(os.path.join(root, file))
            audio = keras.utils.resample(audio, sr, 16000)
            # 统一长度（同步骤1的target_len）
            target_len = 160000
            if len(audio) > target_len:
                audio = audio[:target_len]
            else:
                audio = np.pad(audio, (0, target_len - len(audio)))
            # 保存为numpy数组
            title = os.path.splitext(file)[0]
            np.save(os.path.join(preprocessed_dir, f"{title}.npy"), audio)

# 加载预处理好的数据训练
xs, ys = [], []
for file in os.listdir(preprocessed_dir):
    title = os.path.splitext(file)[0]
    x = np.load(os.path.join(preprocessed_dir, file))
    xs.append(np.expand_dims(x, axis=-1))
    ys.append(metadata.loc[title, "UserLikes"])

xs = np.array(xs)
ys = np.array(ys)

model.fit(xs, ys, epochs=1, batch_size=8, verbose=2)

4. 更换更高效的模型结构 #

直接用RNN处理原始音频时域序列计算量极大，建议将音频转换为梅尔频谱图后用CNN训练，计算效率和内存占用都会大幅降低：

def audio_to_melspec(audio, label):
    # 转换为梅尔频谱图
    melspec = keras.layers.MelSpectrogram(
        sample_rate=16000, n_mels=64, fft_length=1024
    )(audio)
    # 添加通道维度适配CNN
    melspec = tf.expand_dims(melspec, axis=-1)
    return melspec, label

# 处理数据集
data = data.map(audio_to_melspec, num_parallel_calls=tf.data.AUTOTUNE)

# 构建CNN模型
inputs = keras.layers.Input((None, 64, 1))
x = keras.layers.Conv2D(32, (3, 3), activation="relu")(inputs)
x = keras.layers.MaxPooling2D((2, 2))(x)
x = keras.layers.Conv2D(64, (3, 3), activation="relu")(x)
x = keras.layers.MaxPooling2D((2, 2))(x)
x = keras.layers.Flatten()(x)
x = keras.layers.Dense(128, activation="relu")(x)
output = keras.layers.Dense(1)(x)

model = keras.Model(inputs, output)
model.compile(loss="mean_squared_error")

5. 检查硬件加速配置 #

如果使用CPU训练长序列RNN，计算速度会极慢甚至假死。运行以下代码确认是否启用GPU加速：

import tensorflow as tf
print(tf.config.list_physical_devices('GPU'))

如果无GPU设备，建议安装TensorFlow GPU版本并配置对应驱动；若只能用CPU，务必减小模型规模（比如把SimpleRNN单元数从200降到64）和batch_size。

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