TensorFlow自定义损失函数已做值裁剪仍出现NaN/Inf问题求助

核心问题根源 #

你遇到的NaN问题主要来自两个关键错误：

1. 不必要的矩阵裁剪：对correct和wrong整个矩阵做裁剪，把原本为0的元素强制设为epsilon，导致计算L10范数时引入了极小值项，反向传播时梯度剧烈放大，引发参数更新爆炸，最终导致模型输出出现Inf/NaN。
2. 过小的epsilon值：默认的tf.keras.backend.epsilon()（约1e-7）会让correct_score可能非常小，-log(correct_score)的梯度会达到1e7量级，直接冲垮优化过程。

另外，你的model.compile代码存在语法错误（多了一个右括号），虽然不影响运行，但需要修正：

model.compile(optimizer='adam',
              loss=custom,
              metrics=['accuracy'])

修改后的自定义损失函数 #

下面是修复后的损失函数，解决了上述问题并优化了稳定性：

def custom(target, pred):
    k = 10  # 类别数
    # 将标签转为one-hot编码
    target_oh = tf.one_hot(tf.squeeze(target), k)
    
    # 直接提取正确标签的概率（无需裁剪整个矩阵）
    correct_prob = tf.reduce_sum(target_oh * pred, axis=1)
    # 错误标签的概率矩阵（正确标签位置为0）
    wrong_probs = pred * (1 - target_oh)
    
    # 使用更大的epsilon，避免梯度爆炸
    _epsilon = 1e-4
    # 裁剪正确概率，确保log计算安全
    correct_prob = tf.clip_by_value(correct_prob, _epsilon, 1. - _epsilon)
    correct_cost = -tf.math.log(correct_prob)
    
    # 计算错误标签的L10范数
    wrong_norm = tf.norm(wrong_probs, ord=10, axis=1)
    # 裁剪范数，防止数值误差导致超过1
    wrong_norm = tf.clip_by_value(wrong_norm, _epsilon, 1. - _epsilon)
    wrong_score = 1. - wrong_norm
    wrong_score = tf.clip_by_value(wrong_score, _epsilon, 1. - _epsilon)
    wrong_cost = -tf.math.log(wrong_score)
    
    # 对错误项加权，降低其对总损失的影响，提升稳定性
    cost = correct_cost + 0.5 * wrong_cost
    return cost

额外优化建议 #

1. 降低学习率：使用更小的Adam学习率，避免参数更新幅度过大：

   optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
   model.compile(optimizer=optimizer, loss=custom, metrics=['accuracy'])
   

2. 调整L10范数的阶数：如果仍出现NaN，可以尝试降低范数的阶数（比如从10降到5），高阶范数对极小值更敏感，容易引发数值不稳定。

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