TensorFlow中通过回调获取Adam优化器步长的技术咨询

Adam的自适应特性体现在实际权重更新步长Δθ上，而非恒定的基础学习率α。要计算并记录这个步长，关键是获取优化器维护的一阶矩（m）、二阶矩（v）变量，再结合偏差修正规则计算。

1. 获取Adam的一阶矩（m）和二阶矩（v）变量 #

TensorFlow/Keras的Adam优化器会为每个可训练参数单独维护m（一阶滑动平均）和v（二阶滑动平均）变量，你可以通过优化器的get_slot()方法精准获取：

• 先获取模型的所有可训练参数：trainable_params = model.trainable_variables
• 对每个参数，用optimizer.get_slot(param, slot_name)获取对应的m或v，其中slot_name为'm'或'v'

示例代码片段：

import tensorflow as tf

# 假设model是你的训练模型，optimizer是Adam实例
trainable_params = model.trainable_variables
beta1 = optimizer.beta1.numpy()
beta2 = optimizer.beta2.numpy()
epsilon = optimizer.epsilon.numpy()
alpha = optimizer.lr.numpy()
# 获取全局步数（已进行的训练batch数）
t = optimizer.iterations.numpy()

# 计算偏差修正系数
beta1_power = beta1 ** t
beta2_power = beta2 ** t

# 遍历每个参数，计算对应的Δθ
for param in trainable_params:
    # 获取当前参数的m和v
    m = optimizer.get_slot(param, 'm').numpy()
    v = optimizer.get_slot(param, 'v').numpy()
    # 计算偏差修正后的m̂和v̂
    m_hat = m / (1 - beta1_power)
    v_hat = v / (1 - beta2_power)
    # 计算实际步长Δθ
    delta_theta = -alpha * m_hat / (tf.sqrt(v_hat) + epsilon)
    # 这里可以记录delta_theta，比如写入TensorBoard或保存到日志
    print(f"参数{param.name}的步长均值: {tf.reduce_mean(tf.abs(delta_theta))}")

2. 集成到自定义回调（TensorBoard记录） #

你可以把上述逻辑封装到Keras回调中，在每个训练batch结束后计算并记录步长：

class AdamStepCallback(tf.keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        optimizer = self.model.optimizer
        trainable_params = self.model.trainable_variables
        beta1 = optimizer.beta1.numpy()
        beta2 = optimizer.beta2.numpy()
        epsilon = optimizer.epsilon.numpy()
        alpha = optimizer.lr.numpy()
        t = optimizer.iterations.numpy()
        
        beta1_power = beta1 ** t
        beta2_power = beta2 ** t
        
        for param in trainable_params:
            m = optimizer.get_slot(param, 'm').numpy()
            v = optimizer.get_slot(param, 'v').numpy()
            m_hat = m / (1 - beta1_power)
            v_hat = v / (1 - beta2_power)
            delta_theta = -alpha * m_hat / (tf.sqrt(v_hat) + epsilon)
            # 记录步长的均值到TensorBoard
            step_mean = tf.reduce_mean(tf.abs(delta_theta))
            tf.summary.scalar(f"step_size/{param.name}", step_mean, step=t)

使用时，将回调加入模型的fit()方法：

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs')
adam_step_callback = AdamStepCallback()
model.fit(x_train, y_train, epochs=10, callbacks=[tensorboard_callback, adam_step_callback])

关于get_updates()方法的说明 #

get_updates()是优化器内部用来构建参数更新计算图的方法，它返回的是TensorFlow操作张量（而非步长数值），需要传入损失张量和参数列表来生成更新逻辑，并不适合直接用来获取实际步长值，因此不推荐用它来实现你的需求。

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