TensorFlow中损失非零但梯度为0的问题排查求助

问题描述 #

遇到TensorFlow中损失值非零但梯度为0的问题，最小复现代码如下：

import tensorflow as tf
import numpy as np


def setup_mat():
    np_data = np.asarray(
        [
            [-1, -2, -3],
            [-0.1, -0.2, -0.3],
            [0.1, 0.2, 0.3],
            [1, 2, 3]
        ]
    )
    return tf.convert_to_tensor(
        np_data, dtype_hint=tf.float32)


def report_binary_xent():
    Xs = setup_mat()
    weights = tf.Variable([[1.0], [2.0], [3.0]], shape=(3, 1), trainable=True)

    bce = tf.keras.losses.BinaryCrossentropy(from_logits=False)

    labels = np.asarray([0.0, 1.0, 0.0, 0.0])

    with tf.GradientTape(persistent=True) as g:
        g.watch([weights])
        logits = tf.matmul(Xs, weights)
        preds = tf.math.softmax(logits)
        bce_res = bce(labels, preds)
    dydx = g.gradient(bce_res, [preds, logits, weights])
    print(dydx)

认为梯度应该非零的依据：

• weights已标记为可训练（trainable=True）
• 已通过GradientTape监视weights
• 损失值（bce_res）非零

问题原因分析 #

核心问题是二分类场景错误使用了softmax激活函数，具体逻辑如下：

1. 代码中logits的形状是(4,1)，即每个样本对应一个输出值；tf.math.softmax是对张量的最后一维做归一化，单元素维度归一化后结果恒为1.0，因此preds的所有值都是1.0。
2. 二元交叉熵的梯度公式为：(pred - label) / (pred * (1 - pred))，当pred=1.0时，分母1-pred=0，TensorFlow为避免数值异常会将梯度置为0。
3. 损失值非零是因为TensorFlow对log(0)做了数值截断处理（避免直接出现无穷大），但这无法改变梯度为0的结果。

补充：GradientTape默认会自动监视tf.Variable，因此g.watch([weights])属于冗余代码，但不是导致梯度为0的原因。

修正方案 #

针对二分类场景，有两种正确的处理方式：

方案1：使用sigmoid激活函数 + from_logits=False #

将softmax替换为sigmoid，对应二分类场景输出单个样本的正类概率：

import tensorflow as tf
import numpy as np


def setup_mat():
    np_data = np.asarray(
        [
            [-1, -2, -3],
            [-0.1, -0.2, -0.3],
            [0.1, 0.2, 0.3],
            [1, 2, 3]
        ]
    )
    return tf.convert_to_tensor(
        np_data, dtype_hint=tf.float32)


def report_binary_xent():
    Xs = setup_mat()
    weights = tf.Variable([[1.0], [2.0], [3.0]], shape=(3, 1), trainable=True)

    bce = tf.keras.losses.BinaryCrossentropy(from_logits=False)

    labels = np.asarray([0.0, 1.0, 0.0, 0.0])

    with tf.GradientTape(persistent=True) as g:
        logits = tf.matmul(Xs, weights)
        preds = tf.math.sigmoid(logits)  # 替换为sigmoid
        bce_res = bce(labels, preds)
    dydx = g.gradient(bce_res, [preds, logits, weights])
    print("损失值：", bce_res.numpy())
    print("梯度：", [d.numpy() if d is not None else d for d in dydx])

方案2：设置from_logits=True，直接传入logits #

让损失函数内部处理数值稳定的sigmoid计算，无需手动添加激活函数：

import tensorflow as tf
import numpy as np


def setup_mat():
    np_data = np.asarray(
        [
            [-1, -2, -3],
            [-0.1, -0.2, -0.3],
            [0.1, 0.2, 0.3],
            [1, 2, 3]
        ]
    )
    return tf.convert_to_tensor(
        np_data, dtype_hint=tf.float32)


def report_binary_xent():
    Xs = setup_mat()
    weights = tf.Variable([[1.0], [2.0], [3.0]], shape=(3, 1), trainable=True)

    bce = tf.keras.losses.BinaryCrossentropy(from_logits=True)  # 开启from_logits

    labels = np.asarray([0.0, 1.0, 0.0, 0.0])

    with tf.GradientTape(persistent=True) as g:
        logits = tf.matmul(Xs, weights)
        bce_res = bce(labels, logits)  # 直接传入logits
    dydx = g.gradient(bce_res, [logits, weights])
    print("损失值：", bce_res.numpy())
    print("梯度：", [d.numpy() if d is not None else d for d in dydx])

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