带掩码与自定义损失函数的Keras LSTM首轮训练后失效

先修正明显的语法错误 #

你的模型定义中，Dense层的初始化器参数传递存在语法错误，这会导致初始化逻辑失效：

# 错误写法
model.add(tf.keras.layers.Dense(30), 
    kernel_initializer=tf.keras.initializers.zeros())

# 正确写法
model.add(tf.keras.layers.Dense(30, 
    kernel_initializer=tf.keras.initializers.zeros()))

修正后再进行后续排查，避免无关问题干扰。

最可能的核心原因：Masking层使用-inf引发数值不稳定 #

Keras的Masking层虽能标记屏蔽时间步，但-inf本身会在LSTM内部计算中直接触发数值异常：

• LSTM的输入门、遗忘门依赖sigmoid激活，-inf输入会直接输出0，后续状态更新（如c_t = f_t * c_{t-1} + i_t * tanh(z_t)）会因-inf的参与产生NaN。
• 即使Masking层标记了屏蔽状态，序列开头的-inf填充仍可能影响LSTM的初始状态计算，首轮训练后权重被污染，导致后续批次全量NaN。

解决方法：
将掩码值替换为有限的、真实数据中不会出现的数值，比如-1e9（需确保真实数据无此值）：

model.add(tf.keras.layers.Masking(mask_value=-1e9, 
    input_shape=(train_X.shape[1], train_X.shape[2])))

同时检查真实输入数据是否混入异常值：

print("训练数据是否含非有限值:", tf.math.is_finite(train_X).numpy().all())

自定义损失函数的梯度排查 #

你提到损失输出为样本级1D张量，这符合Keras要求，但需确认梯度是否存在NaN——即使损失值正常，梯度NaN会导致优化器更新权重时产生NaN，进而破坏后续计算。

排查代码：

import tensorflow as tf

# 取一个批次的样本测试
batch_x = train_X[:32]
batch_y = train_y[:32]

with tf.GradientTape() as tape:
    preds = model(batch_x, training=True)
    loss = batched_custom_loss(batch_y, preds)

# 计算所有可训练参数的梯度并检查NaN
grads = tape.gradient(loss, model.trainable_variables)
for idx, grad in enumerate(grads):
    var_name = model.trainable_variables[idx].name
    has_nan = tf.math.is_nan(grad).numpy().any()
    print(f"参数 {var_name} 的梯度是否含NaN: {has_nan}")

若发现梯度NaN，需修正损失函数的数值稳定性：

• 避免除以0：将x / y改为x / (y + 1e-8)
• 避免对数输入非正：将tf.log(x)改为tf.log(tf.maximum(x, 1e-8))
• 避免平方根输入负数：将tf.sqrt(x)改为tf.sqrt(tf.maximum(x, 0.))

其他辅助排查步骤 #

1. 测试标准损失函数：暂时将自定义损失替换为MSE等标准损失，若训练不再出现NaN，说明问题出在自定义损失；若仍出现NaN，问题则在模型或数据层面。
2. 检查LSTM层输出：手动前向传播多个批次，确认LSTM输出是否异常：

# 构建输出LSTM层结果的模型
lstm_output_model = tf.keras.Model(inputs=model.input, outputs=model.layers[1].output)

batch1_output = lstm_output_model(train_X[:32])
batch2_output = lstm_output_model(train_X[32:64])

print("批次1 LSTM输出是否含NaN:", tf.math.is_nan(batch1_output).numpy().any())
print("批次2 LSTM输出是否含NaN:", tf.math.is_nan(batch2_output).numpy().any())

若批次2输出已含NaN，说明掩码或输入数据的问题导致LSTM计算异常。

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