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 做线性上升段。
更稳妥的做法是手写一个两段式调度：
立即学习“Python免费学习笔记（深入）”；
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 < self.warmup_steps,
lambda: self.init_lr + (self.peak_lr - self.init_lr) * (step / self.warmup_steps),
lambda: tf.cond(
step < self.decay_start_step,
lambda: self.peak_lr,
lambda: self.peak_lr * (self.decay_rate ** ((step - self.decay_start_step) // (1 if self.staircase else 1)))
)
)
return lr
注意：这里用了 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 < 5:
return 1e-5 + (1e-3 - 1e-5) * epoch / 5 # warmup for 5 epochs
else:
return 1e-3 * 0.5 * (1 + np.cos(np.pi * (epoch - 5) / (50 - 5)))
<p>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 就够用，而且不容易出兼容性问题。
