Keras自定义加权交叉熵损失训练时出现NaN问题求助

问题分析 #

训练过程中第5-10轮突然出现NaN损失，错误提示显示模型输出存在NaN值。结合代码来看，问题根源在于自定义损失函数的数值不稳定操作，以及可能的梯度爆炸导致模型参数更新异常，进而让softmax输出出现NaN。

解决方案 #

1. 修复自定义损失函数的数值问题 #

你的损失函数中手动对y_pred做归一化（y_pred /= K.sum(y_pred[...,:], axis=-1, keepdims=True)）是多余的——模型输出已经经过softmax激活，本身就满足类别概率和为1。这一步操作可能因数值下溢导致sum为0，进而引发NaN。同时调整损失的求和逻辑，避免维度错误：

def weighted_categorical_crossentropy(weights=the_weights):
    weights = K.variable(weights)
        
    def loss(y_true, y_pred):
        # 裁剪预测值，避免log(0)或log(1)导致的数值异常
        y_pred = K.clip(y_pred, K.epsilon(), 1 - K.epsilon())
        # 计算加权交叉熵，确保维度匹配
        cross_entropy = y_true * K.log(y_pred) * weights
        # 按类别维度求和后取样本平均，提升数值稳定性
        loss = -K.mean(K.sum(cross_entropy, axis=-1))
        return loss
    
    return loss

2. 添加梯度裁剪防止梯度爆炸 #

梯度爆炸会导致模型参数更新后出现极端值，进而让softmax输出NaN。在优化器中添加梯度裁剪限制：

optimizer_ = tf.keras.optimizers.Adam(learning_rate=lr_rate, clipnorm=1.0)
# 或使用clipvalue限制单个梯度的绝对值：clipvalue=0.5

3. 优化模型层的初始化与配置 #

• 给卷积层添加He初始化（适合ReLU/LeakyReLU激活），提升初始数值稳定性：

x = tf.keras.layers.Conv2D(64, (4, 4), strides=(1,1), use_bias=False, padding='same', 
                           name='Conv1', kernel_initializer='he_normal')(inp)
x = tf.keras.layers.Conv2D(32, (3, 3), strides=(1,1), use_bias=False, padding='same', 
                           name='Conv2', kernel_initializer='he_normal')(x)

• 确保BatchNormalization在训练时正常更新：训练时不要手动设置training=False，Keras会自动处理训练/推理模式。

4. 降低学习率 #

原学习率0.001可能过高，尝试降低到0.0001或更小，避免参数更新幅度过大导致数值异常：

def make_model(lr_rate=0.0001):  # 调整初始学习率
    # ... 模型结构不变

5. 验证指标的兼容性 #

你使用的部分指标（如AUC）在高维输出（(S1,S2,13)）下可能需要调整参数。可以暂时移除部分指标（如AUC），先保证训练稳定，再逐步添加验证：

metrics=[
    keras.metrics.TruePositives(name='tp'),
    keras.metrics.FalsePositives(name='fp'),
    keras.metrics.Precision(name='precision'),
    keras.metrics.Recall(name='recall')
]

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