PyTorch编解码器TTS模型训练中损失爆炸问题求助

我基于LibriSpeech 100小时数据集，使用PyTorch搭建编码器-解码器架构的文本转语音（TTS）模型，输入文本输出梅尔频谱图。但无论设置何种批次大小，训练进行到第2至第3批时损失都会出现爆炸情况，从初始的0.2左右飙升至1e17量级，最终变为inf。怀疑训练代码存在问题，但无法准确定位，恳请提供排查思路。

训练时终端输出如下：

loss tensor(0.2049, device='cuda:0', grad_fn=<MSELossBackward0>)
loss tensor(8.0709e+17, device='cuda:0', grad_fn=<MSELossBackward0>)
loss tensor(9.1250e+29, device='cuda:0', grad_fn=<MSELossBackward0>)
... [loss will be inf]

排查思路 #

• 数据预处理校验：确认梅尔频谱图的归一化/标准化是否正确。梅尔谱的数值范围通常需要缩放到[-1,1]或[0,1]区间，如果目标数据未做处理，模型输出与目标的数值差会引发MSE损失急剧放大。同时检查文本输入的编码（如音素转索引）是否存在异常值。
• 模型参数初始化检查：编码器、解码器的线性层、注意力层初始化是否合理。默认初始化可能导致输出值过大，建议改用Xavier或He初始化；注意力权重若初始化不当，易产生极端权重触发梯度爆炸。
• 添加梯度裁剪：如果训练代码未实现梯度裁剪，立刻加入torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)，限制梯度的最大范数，避免梯度突变导致权重异常更新。
• 调整学习率：初始学习率可能过高，TTS模型通常适合1e-4~1e-5量级的学习率，先尝试降低学习率观察效果。同时检查优化器参数（如Adam的beta值）是否需要调整。
• 模型结构合理性检查：解码器输出层是否添加了合适的激活函数？梅尔谱的数值范围有限，若输出层未用Sigmoid/Tanh限制范围，模型可能输出极大值引发损失爆炸。另外检查注意力机制实现，是否存在softmax输入过大导致数值溢出的情况。
• 数值稳定性排查：训练时打印模型输出、目标梅尔谱的数值范围，确认是否存在异常大的样本；使用torch.autograd.detect_anomaly()定位反向传播中出现NaN/inf的层，精准找到问题来源。
• 更换损失函数：MSE对异常值敏感度高，若数据存在异常样本，可尝试改用L1损失，验证是否仍会出现损失爆炸。

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