TensorFlow 2.x 中动态调整学习率的API主要包括:tf.keras.optimizers.schedules.ExponentialDecay、PiecewiseConstantDecay、PolynomialDecay等LearningRateSchedule子类,以及tf.keras.callbacks.LearningRateScheduler回调;前者基于step自动计算,后者按epoch调用Python函数,二者均被官方推荐且兼容分布式训练与模型保存。

TensorFlow里哪些API能动态改学习率
TensorFlow 2.x 默认使用 tf.keras.optimizers.Optimizer(如 Adam、SGD),它的学习率是可变的张量(tf.Variable),所以只要在训练过程中更新它,就能实现动态调整。核心不是换优化器,而是控制 optimizer.learning_rate 这个变量。
官方推荐方式是用学习率调度器(LearningRateSchedule),比如 tf.keras.optimizers.schedules.ExponentialDecay、PiecewiseConstantDecay、PolynomialDecay 等——它们返回的是一个 callable,Keras 在每步训练时自动调用它算出当前学习率,无需手动干预。
-
ExponentialDecay适合平滑下降,但容易后期过慢; -
PiecewiseConstantDecay适合按 epoch 阶梯式下调(比如 ImageNet 训练常用); -
PolynomialDecay支持 warmup + decay 组合,更贴近 LAMB 或 AdamW 的实践; - 自定义函数必须继承
tf.keras.optimizers.schedules.LearningRateSchedule并实现__call__和get_config,否则无法保存模型。
如何用ExponentialDecay做warmup+decay混合调度
纯 ExponentialDecay 不支持 warmup,但可以用 tf.keras.optimizers.schedules.PolynomialDecay 模拟:设初始学习率为 0,power=1.0,end_learning_rate 为最终值,再配合 initial_learning_rate 做线性上升段。
更稳妥的做法是手写一个两段式调度:
class WarmupThenDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
def __init__(self, init_lr, peak_lr, warmup_steps, decay_start_step, decay_rate, staircase=False):
self.init_lr = init_lr
self.peak_lr = peak_lr
self.warmup_steps = tf.cast(warmup_steps, tf.float32)
self.decay_start_step = decay_start_step
self.decay_rate = decay_rate
self.staircase = staircase
<pre class="brush:php;toolbar:false;">def __call__(self, step):
step = tf.cast(step, tf.float32)
lr = tf.cond(
step <p></p>
注意:这里用了 tf.cond 而非 Python if,因为要支持图模式;step 是全局训练步数(不是 epoch),来自 tf.keras.callbacks.Callback 的 logs['step'] 或 optimizer.iterations。
训练循环中手动修改learning_rate会有什么问题
直接赋值 optimizer.learning_rate.assign(new_lr) 是可行的,但容易踩三个坑:
- 如果用
model.fit(),Keras 内部可能缓存了学习率值(尤其开启tf.function图编译后),导致赋值不生效; - 多 GPU 分布式训练(
tf.distribute.MirroredStrategy)下,assign必须在strategy.run()内或使用strategy.experimental_variable_creator_scope()创建变量,否则只改了副本; - 保存模型时,手动 assign 的值不会被自动序列化进
optimizer.get_weights(),加载后学习率会回到初始值; - 用
tf.keras.callbacks.LearningRateScheduler回调比手动 assign 更安全,它在每个 epoch 开始前调用,且兼容分布式和保存/加载。
为什么用LearningRateScheduler回调比自定义schedule更省心
tf.keras.callbacks.LearningRateScheduler 接收一个 Python 函数,输入是 epoch(从 0 开始),输出是学习率标量。它绕过了图构建限制,调试友好,也不依赖 step 计数逻辑。
例如实现 cosine annealing with warmup:
def scheduler(epoch):
if epoch lr_callback = tf.keras.callbacks.LearningRateScheduler(scheduler)
model.fit(..., callbacks=[lr_callback])
这个写法简单直观,但注意:它只在 epoch 级别生效,不适合 step-level 的精细控制(比如 BERT 的 1000 步 warmup);另外 NumPy 函数不能直接用于 tf.function,所以不要在自定义训练循环里混用。
真正需要 step 级调度时,优先选 LearningRateSchedule 子类;只需要按 epoch 调整,LearningRateScheduler 就够用,而且不容易出兼容性问题。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











