TensorFlow Eager Execution结合LR decay失效问题求助

嘿，我完全懂你这种卡在Eager Execution下学习率衰减不生效的挫败感——一开始我也怀疑是不是TensorFlow的bug，但其实大多是因为Eager模式下的张量更新逻辑和我们习惯的Graph模式不一样。咱们来一步步拆解问题，找到解决方案！

核心问题所在 #

在Eager模式下，静态张量是无法被动态更新的。如果你的学习率衰减逻辑只是用tf.constant定义了一个初始值，或者没有把衰减步骤和训练迭代绑定，就会出现“学习率张量看起来完全没变化”的情况。

两种可行的解决方案 #

方案1：用TensorFlow内置的学习率调度器（推荐）

TensorFlow的tf.keras.optimizers.schedules模块提供了多种开箱即用的衰减策略，Eager模式下会自动追踪训练步数并更新学习率，不需要手动维护状态。比如指数衰减、分段常数衰减等。

以指数衰减为例，直接把调度器传给优化器即可：

import tensorflow as tf
NUM_SAMPLED = 64

# 定义学习率衰减策略
initial_learning_rate = 0.01
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate,
    decay_steps=5000,  # 每5000步衰减一次
    decay_rate=0.9,    # 每次衰减为当前的90%
    staircase=True     # 阶梯式衰减（False则是连续衰减）
)

# 初始化优化器时传入调度器
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)

在训练循环里，你只需要正常调用optimizer.apply_gradients，Eager模式会自动计算当前的学习率，完全不需要额外操作。

方案2：手动维护学习率变量（适合自定义衰减逻辑）

如果你的衰减逻辑比较特殊（比如自定义的非线性衰减），可以用tf.Variable来保存当前学习率，然后在每步训练后手动更新它：

# 初始化学习率为可训练变量（注意trainable=False，避免被优化器更新）
current_lr = tf.Variable(initial_value=0.01, trainable=False)

# 自定义衰减函数
def update_learning_rate(step):
    # 比如每1000步学习率减半
    if step % 1000 == 0 and step != 0:
        current_lr.assign(current_lr * 0.5)
    return current_lr

# 训练循环中
for epoch in range(10):
    for step, (x, y) in enumerate(train_dataset):
        # 先更新学习率
        update_learning_rate(step)
        # 把更新后的学习率传给优化器
        optimizer.lr.assign(current_lr)
        
        # 正常计算梯度并更新参数
        with tf.GradientTape() as tape:
            loss = model(x, y)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

常见错误排查 #

• 不要用静态张量定义学习率：比如lr = tf.constant(0.01)，这种张量是不可变的，每次训练都会用初始值，自然看不到衰减效果。
• 确认衰减逻辑被触发：检查decay_steps是不是设置得太大，导致还没到衰减的步数；或者有没有在每步训练后调用更新学习率的代码。
• 用fit方法时要正确传入调度器：如果用model.fit而不是自定义训练循环，一定要把带调度器的优化器传入optimizer参数，而不是单独定义学习率。

完整的Word2Vec示例片段 #

补全你的代码，结合内置调度器的完整示例：

import tensorflow as tf
NUM_SAMPLED = 64

class Word2Vec(tf.keras.Model):
    def __init__(self, vocab_size, embed_size, num_sampled=NUM_SAMPLED):
        super().__init__()
        self.vocab_size = vocab_size
        self.num_sampled = num_sampled
        self.embedding = tf.keras.layers.Embedding(vocab_size, embed_size)
        # NCE损失的权重和偏置定义为可训练变量
        self.nce_weights = tf.Variable(tf.random.normal([vocab_size, embed_size]))
        self.nce_biases = tf.Variable(tf.zeros([vocab_size]))

    def call(self, inputs, labels=None):
        embeds = self.embedding(inputs)
        if labels is None:
            return embeds
        # 计算NCE损失
        loss = tf.reduce_mean(tf.nn.nce_loss(
            weights=self.nce_weights,
            biases=self.nce_biases,
            labels=labels,
            inputs=embeds,
            num_sampled=self.num_sampled,
            num_classes=self.vocab_size))
        return loss

# 初始化模型
vocab_size = 10000
embed_size = 128
model = Word2Vec(vocab_size, embed_size)

# 配置学习率衰减和优化器
initial_lr = 0.01
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_lr, decay_steps=5000, decay_rate=0.9, staircase=True)
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)

# 模拟训练数据（实际替换成你的数据集）
train_inputs = tf.random.uniform((100000,), 0, vocab_size, dtype=tf.int32)
train_labels = tf.expand_dims(tf.random.uniform((100000,), 0, vocab_size, dtype=tf.int32), axis=1)
train_dataset = tf.data.Dataset.from_tensor_slices((train_inputs, train_labels)).batch(32)

# Eager模式训练循环
for epoch in range(5):
    total_loss = 0.0
    for step, (x, y) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            loss = model(x, y)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        
        # 每100步打印当前学习率和损失，验证衰减是否生效
        if step % 100 == 0:
            current_lr = optimizer.lr.numpy()
            print(f"Epoch {epoch+1}, Step {step}, Loss: {loss.numpy():.4f}, Current LR: {current_lr:.6f}")
        total_loss += loss.numpy()
    print(f"=== Epoch {epoch+1} Done, Average Loss: {total_loss/len(train_dataset):.4f} ===\n")

运行这个代码，你会看到每经过5000步，学习率就会下降到当前的90%，完全符合预期。

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