GAN训练中tape.gradient返回None值的问题排查求助

问题描述 #

训练GAN时，调用tape.gradient(loss, model.trainable_variables)计算损失梯度，结果返回全为None的数组。检查loss和model.trainable_variables均已正常定义，相关代码如下：

导入依赖 #

import numpy as np
import tensorflow as tf

生成器模型 #

def make_model():
  Model = tf.keras.models.Sequential()
  Model.add(tf.keras.layers.Dense(units = 6, activation='relu'))
  Model.add(tf.keras.layers.Dense(units = 2, activation='relu'))
  Model.add(tf.keras.layers.Dense(units = 2, activation='relu'))
  return Model

测试张量 #

tensor_one = tf.constant([[4, 5]])
tensor_two  = tf.constant([[9, 3]])

判别器模型 #

def make_discriminator_model():
    Model = tf.keras.models.Sequential()
    Model.add(tf.keras.layers.Dense(units = 6, activation="relu"))
    Model.add(tf.keras.layers.Dense(units = 2, activation="relu"))
    Model.add(tf.keras.layers.Dense(units=1, activation='sigmoid'))
    return Model

模型实例化 #

basic_model = make_model()
basic_model(tensor_one)

discriminator_model = make_discriminator_model()
discriminator_model(tensor_one)

损失函数定义 #

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def model_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)

def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss

优化器与训练循环 #

generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

epochs = 100
for epoch in range(epochs):
    with tf.GradientTape() as model_tape, tf.GradientTape() as discrim_tape:
        output_values = basic_model(tensor_one, training=True)
        real_output = discriminator_model(tensor_two, training=True)
        fake_output = discriminator_model(output_values, training=True)

        gen_loss = model_loss(fake_output)
        disc_loss = discriminator_loss(real_output, fake_output)
    
    gradients_of_model = model_tape.gradient(gen_loss, basic_model.trainable_variables)
    gradients_of_discriminator = discrim_tape.gradient(disc_loss, discriminator_model.trainable_variables)

    print(gradients_of_model)

    generator_optimizer.apply_gradients(zip(gradients_of_model, basic_model.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator_model.trainable_variables))

问题根源 #

核心矛盾在于判别器输出与损失函数的参数不匹配：

• 判别器最后一层使用了sigmoid激活，输出的是0-1之间的概率值；
• 但损失函数BinaryCrossentropy设置了from_logits=True，该参数要求输入是未经过激活的原始logits（即网络最后一层的原始输出）。

这种不匹配会导致损失计算时的数值异常，打断梯度回传链路，最终使得tape.gradient无法计算出有效梯度，返回全None。

解决方法 #

有两种等价的修正方案，推荐第一种以避免数值不稳定：

方案1：移除判别器的sigmoid激活（推荐） #

保持损失函数的from_logits=True，直接让判别器输出logits：

def make_discriminator_model():
    Model = tf.keras.models.Sequential()
    Model.add(tf.keras.layers.Dense(units=6, activation="relu"))
    Model.add(tf.keras.layers.Dense(units=2, activation="relu"))
    Model.add(tf.keras.layers.Dense(units=1))  # 移除sigmoid激活
    return Model

方案2：修改损失函数的from_logits参数 #

如果要保留判别器的sigmoid激活，将损失函数的from_logits改为False：

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

验证 #

修正后重新运行训练循环，gradients_of_model和gradients_of_discriminator将返回正常的梯度张量，不再是全None。

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