使用tf.keras.backend.gradient计算梯度结果全为0的原因排查

问题描述 #

使用Keras后端计算MNIST训练模型输入的梯度时，得到的结果始终全为0，即使loss值显示正常（如1e-07）。相关代码如下：

from tensorflow.keras import backend as K, losses, Model, models
import tensorflow as tf
import numpy as np
tf.compat.v1.disable_eager_execution()

# pre-operation
model = models.load_model("")
image = ""
y_true = ""


def fun(model, image, y_true):
    y_pred = model.output

    loss = losses.categorical_crossentropy(y_true, y_pred)

    gradient = K.gradients(loss, model.input)
    gradient = gradient[0]

    sess = tf.compat.v1.keras.backend.get_session()
    res = sess.run(gradient, feed_dict={model.input: image})
    if res.any() == 0.:
        print('all is 0')

    return res

img_attack = fun(model, image, y_true)

梯度输出示例：

...

  [[0.]
   [0.]
   [0.]
   ...
   [0.]
   [0.]
   [0.]]

  [[0.]
   [0.]
   [0.]
   ...
   [0.]
   [0.]
   [0.]]

  [[0.]
   [0.]
   [0.]
   ...
   [0.]
   [0.]
   [0.]]]]

已确认loss、y_pred等参数正常，但梯度全为0，怀疑是softmax后y_pred与y_true过于接近，但无法解释梯度完全为0的情况。

可能的原因及解决办法 #

1. 数值精度截断导致梯度被置0 #

当y_pred与y_true极度接近时，交叉熵损失的梯度会极小，可能低于float32的数值精度阈值，最终被显示为0。可以通过以下方式验证：

• 给输入或标签添加微小扰动，打破完全匹配的状态：

  # 给y_true加扰动
  y_true = y_true + 1e-6
  # 或者给image加扰动
  image = image + np.random.normal(0, 1e-6, image.shape)
  

• 改用float64数据类型计算梯度，提升精度：

  image = image.astype(np.float64)
  model = model.astype('float64')
  

2. 禁用Eager Execution导致计算图异常 #

在TensorFlow 2.x环境下，使用tf.compat.v1.disable_eager_execution()可能会引发计算图构建的隐性问题，建议改用TF2.x原生的Eager模式，用tf.GradientTape计算梯度，代码示例：

from tensorflow.keras import losses, Model, models
import tensorflow as tf
import numpy as np

# 加载模型和数据（确保image是batch维度，如(1,28,28,1)）
model = models.load_model("your_model_path.h5")
image = np.expand_dims(your_mnist_image, axis=0)
y_true = np.expand_dims(your_one_hot_label, axis=0)

def compute_gradient(model, image, y_true):
    image_tensor = tf.convert_to_tensor(image, dtype=tf.float32)
    with tf.GradientTape() as tape:
        tape.watch(image_tensor)
        # training=False避免影响BatchNorm、Dropout等层的行为
        y_pred = model(image_tensor, training=False)
        loss = losses.categorical_crossentropy(y_true, y_pred)
    gradient = tape.gradient(loss, image_tensor)
    return gradient.numpy()

img_attack = compute_gradient(model, image, y_true)
if (img_attack == 0).all():
    print('all is 0')

3. 模型输出层与损失函数的匹配问题 #

如果模型输出层是softmax，categorical_crossentropy在Keras内部会做合并优化，但手动计算梯度时可能存在数值稳定性问题。可以尝试：

• 如果模型能输出logits（未经过softmax的原始输出），修改损失函数使用from_logits=True参数：

  # 假设模型有一个输出logits的层，比如命名为'logits'
  logits_model = Model(inputs=model.input, outputs=model.get_layer('logits').output)
  y_pred = logits_model(image_tensor)
  loss = losses.categorical_crossentropy(y_true, y_pred, from_logits=True)
  

  这样可以避免softmax后的数值饱和问题，提升梯度计算的稳定性。

4. 模型层导致梯度消失 #

检查模型中是否存在容易引发梯度消失的层：

• 若使用了ReLU激活函数，当前输入可能落在ReLU的负区间，导致梯度为0。可以尝试替换为LeakyReLU，或者检查模型各层的激活值是否处于饱和状态。
• 若模型深度较深，可能存在梯度消失问题，可以尝试添加残差连接等结构。

5. 输入数据维度/类型问题 #

• 确保image是正确的batch格式（MNIST应为(batch_size, 28, 28, 1)），y_true是one-hot编码且维度与模型输出一致。
• 检查image的数据类型是否为float32，整数类型无法计算梯度。

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