如何在Keras中通过Weights & Biases记录学习率调度器的数值变化

WandbCallback默认不支持自动采集tf.keras.optimizers.schedules.LearningRateSchedule类型调度器的实时学习率，可通过以下两种方式实现记录：

场景1：使用model.fit训练 #

自定义Keras回调，每轮/每步自动拉取当前学习率上报到Wandb：

1. 先定义回调类

import wandb
import tensorflow as tf
from tensorflow import keras as tfk

class LearningRateLogger(tfk.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        # 调度器会根据优化器的当前迭代步数自动计算对应学习率
        current_lr = self.model.optimizer.lr(self.model.optimizer.iterations)
        # 上报到Wandb，step和训练轮次对齐
        wandb.log({"learning_rate": current_lr.numpy()}, step=epoch)

2. 将自定义回调和WandbCallback一同传入model.fit的回调列表即可

train_callbacks = [
    wandb.keras.WandbCallback(),
    LearningRateLogger()
]

model.fit(
    train_x, train_y,
    epochs=20,
    callbacks=train_callbacks
)

场景2：使用自定义训练循环 #

直接在迭代逻辑中插入学习率记录代码即可：

for epoch in range(total_epochs):
    for step, (batch_x, batch_y) in enumerate(train_loader):
        # 原有前向传播、梯度更新逻辑
        with tf.GradientTape() as tape:
            pred = model(batch_x, training=True)
            loss = loss_func(batch_y, pred)
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        # 记录当前学习率，可根据需要选择按步或按轮记录
        current_lr = optimizer.lr(optimizer.iterations)
        wandb.log({
            "train_loss": loss.numpy(),
            "learning_rate": current_lr.numpy()
        })

注意：如果你的学习率是按训练步而非轮次更新，将回调的触发方法从on_epoch_end改为on_train_batch_end即可，同时调整wandb.log的step参数和全局训练步对齐，避免日志时序错乱。

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