TensorFlow中应继承tf.keras.optimizers.schedules.LearningRateSchedule实现__call__(self, step)，用tf.*操作处理tensor输入，避免Python原生函数和if分支；warmup+cosine需统一全局step计数，total_steps设为总训练步数，确保checkpoint兼容与相位对齐。
TensorFlow里怎么用自定义函数控制学习率变化
直接用 tf.keras.optimizers.schedules.LearningRateSchedule 子类是最稳妥的方式，比在训练循环里手动改 optimizer.learning_rate 更可靠——后者容易被 checkpoint 保存/恢复逻辑干扰，也绕过梯度更新前的 learning rate 计算时机。
常见错误是写个普通 Python 函数返回 float，然后传给 optimizer，结果报错 TypeError: 'float' object is not callable，因为 Keras 期望的是 callable 对象，且必须支持 tensor 输入（训练步数 step 是 tf.Tensor 类型）。
- 继承
LearningRateSchedule，实现__call__(self, step)方法，输入step是 int64 tensor，返回标量 float32 tensor - 所有计算必须用
tf.*操作（如tf.math.maximum、tf.cast），不能用max()、int()等原生 Python 函数 - 若需多段逻辑（比如 warmup → decay → hold），用
tf.cond或tf.case，避免 Python if 分支（会静态绑定，无法随 step 动态切换）
多步长训练中 warmup + cosine decay 怎么写
这是最常被复制粘贴但又最容易出错的组合：warmup 阶段学习率线性上升，之后接余弦退火。关键陷阱在于 step 计数是否包含 warmup 阶段——很多实现把 warmup 步数单独计数，导致总 step 偏移，decay 周期错位。
正确做法是统一用全局 step，用条件判断分段：
立即学习“Python免费学习笔记（深入）”；
class WarmupCosineDecay(LearningRateSchedule):
def __init__(self, init_lr, peak_lr, warmup_steps, total_steps):
self.init_lr = init_lr
self.peak_lr = peak_lr
self.warmup_steps = tf.cast(warmup_steps, tf.float32)
self.total_steps = tf.cast(total_steps, tf.float32)
<pre class='brush:python;toolbar:false;'>def __call__(self, step):
step = tf.cast(step, tf.float32)
# warmup 阶段：线性从 init_lr 到 peak_lr
lr1 = self.init_lr + (self.peak_lr - self.init_lr) * step / self.warmup_steps
# decay 阶段：cosine 从 peak_lr 到 init_lr
progress = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
lr2 = self.init_lr + 0.5 * (self.peak_lr - self.init_lr) * (1 + tf.cos(np.pi * progress))
# 合并：step < warmup_steps 时取 lr1，否则取 lr2
return tf.where(step < self.warmup_steps, lr1, lr2)注意 total_steps 应设为整个训练的总 step 数（不是 epoch 数），否则 cosine 周期会压缩或拉伸；如果训练中途断点续训，step 会自动从 checkpoint 中恢复，无需额外处理。
为什么用 tf.keras.callbacks.LearningRateScheduler 会失效
这个 callback 只在每个 epoch 开始时调用一次，传入的是 epoch 编号（int），不是 step。它根本无法支撑 step-level 的精细调度（比如每 100 步微调一次 lr），更无法用于 steps_per_epoch=None 的 dataset 流式训练场景。
- 只适用于 epoch 级粗粒度调整，比如每 10 个 epoch 把 lr ×0.9
- 回调中返回的值会被直接赋给
optimizer.learning_rate，但该赋值可能被后续 eager 执行中的 optimizer 内部逻辑覆盖（尤其在混合精度训练中） - 不参与 SavedModel 导出，模型 reload 后 lr 调度逻辑丢失
自定义调度器如何和 checkpoint 兼容
只要调度器本身不持有可训练变量（即没用 tf.Variable），就天然兼容 tf.train.Checkpoint。Keras 在保存 optimizer 时，会自动记录当前 step 和学习率值，但不会保存调度器的 Python 逻辑——所以你必须确保调度器类定义在加载时依然可用，且参数（如 warmup_steps）与保存时一致。
容易忽略的一点：如果你在调度器里用了外部闭包变量（比如从 closure 捕获一个 list），那 checkpoint 加载后会出错，因为闭包状态无法序列化。务必把所有配置都作为 __init__ 参数传入，并存为实例属性。
多步长训练真正难的不是写公式，而是让 step 计数、warmup 边界、decay 周期、checkpoint 恢复这四者对齐——差一个 step，cosine 就可能跳相位，warmup 就可能漏半步。别信“差不多”，全得用 tf.print 在前几 step 打印真实 lr 值来验证。
