tf.keras.callbacks.ModelCheckpoint不能直接续训,因其默认仅保存模型权重,不保存优化器状态、学习率调度器和当前epoch等训练上下文。

为什么 tf.keras.callbacks.ModelCheckpoint 不能直接续训?
很多人以为只要保存了模型文件,加载后就能接着训练——但实际会发现 optimizer 的状态(比如 Adam 的 momentum、velocity)和当前 epoch 数全丢了。模型权重只是训练结果的一部分,ModelCheckpoint 默认只存权重(save_weights_only=True),不保存优化器状态、学习率调度器、甚至随机数生成器的 seed。所以即使模型结构一致,从检查点加载后调用 model.fit(),还是会从 epoch 0 开始,且优化器从头初始化。
必须同时保存和恢复这三样东西
断点续训本质是恢复整个训练上下文。缺一不可:
-
model.save_weights()或完整模型(含结构)——用于恢复网络参数 -
optimizer.get_weights()和optimizer.set_weights()—— 恢复动量、累积梯度等内部状态 - 手动记录并恢复
epoch、step、learning_rate当前值(尤其用了tf.keras.optimizers.schedules时)
推荐用 tf.train.Checkpoint 统一管理,它能原子化保存/恢复多个对象:
checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
checkpoint.save('./ckpt/model') # 生成 model-1.index / .data 等
# 恢复时:
checkpoint.restore(tf.train.latest_checkpoint('./ckpt/'))
使用 tf.keras.callbacks.BackupAndRestore 的前提和坑
这是 Keras 官方提供的“开箱即用”续训回调,但它有硬性限制:
- 只支持
model.fit(),不支持自定义训练循环(tf.GradientTape) - 备份路径必须可写,且不能跨机器/不同 TensorFlow 版本混用(.index 文件格式可能不兼容)
- 默认每 epoch 备份一次,若训练中崩溃在 epoch 中间,会丢失该 epoch 的部分进度(它不支持 step 级备份)
- 恢复时自动跳过已训练的 epoch,但不会重置
tf.data.Dataset的 shuffle seed —— 如果你依赖固定 shuffle 顺序做 debug,这点容易被忽略
启用方式很简单:
backup_cb = tf.keras.callbacks.BackupAndRestore(backup_dir='./backup') model.fit(x_train, y_train, callbacks=[backup_cb])
自定义训练循环中如何安全续训?
如果你用 tf.GradientTape 写训练步,BackupAndRestore 就失效了,必须手管。关键点:
- 用
tf.train.Checkpoint保存model、optimizer、epoch(作为tf.Variable)、甚至train_ds_iterator(如果需要精确恢复数据流位置) - 恢复后,手动把
epoch变量加 1,再传给range(epoch.numpy(), total_epochs) - 注意
tf.data.Dataset的状态无法直接序列化,若需恢复确切 batch 位置,得用dataset.skip(n)或从 checkpoint 里存一个 step 计数器
最容易漏的是:恢复 optimizer.learning_rate。如果是 tf.keras.optimizers.schedules.ExponentialDecay,它的 step 是内部变量,必须显式恢复或重建 scheduler 并设初始 initial_learning_rate。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











