Question
BackupAndRestore omits custom callback state and learning-rate schedules
Cig Atlas
0 reputation · 26 Feb 2021, 17:54 UTC
39K views1
Loading…
0 reputation · 26 Feb 2021, 17:54 UTC
The tf.keras.callbacks.BackupAndRestore callback (stable since TensorFlow 2.6) saves model weights, optimizer state, and the epoch counter at each save_freq interval. Restoration loads these artifacts automatically when model.fit() restarts, resuming training from the next epoch.
Custom callback state, learning-rate schedules, and metric accumulators are not included in the backup. After recovery, any callback that maintains internal counters (e.g., early-stopping patience, learning-rate schedulers with step-dependent logic, or custom metrics that accumulate across batches) resets to its initial state. This can change the training trajectory silently compared to an uninterrupted run.
Given that the callback does not serialize arbitrary callback objects, what patterns exist for reconstructing or persisting this missing state without modifying the core callback? Specifically:
35275 reputation · 26 Feb 2021, 21:34 UTC
The BackupAndRestore callback, stable since TensorFlow 2.6, is designed to persist only model weights, optimizer weights and the epoch counter. It deliberately excludes arbitrary callback internals such as learning‑rate schedule step counters, early‑stopping patience, or metric accumulators.
keras
tensorflowweights, optimizer_weights and epoch to the checkpoint directory.on_epoch_end, writes the relevant attribute of each target callback (e.g., a scheduler’s step counter) to a JSON file alongside the model checkpoint.on_train_begin hook, read that JSON and restore the saved values before model.fit resumes.Please provide: the specific callbacks you need to preserve and whether they expose a serializable attribute (e.g., a numeric counter) that can be saved to a side file.
Use comments to ask for clarification. Post a solution as an answer.
35,275 reputation · 26 Feb 2021, 19:31 UTC
A useful distinction is that a learning-rate schedule implemented as a function of optimizer.iterations may resume correctly in TensorFlow 2.x, because the optimizer state is part of the backup. Schedules whose progress lives inside a callback—such as a custom step counter, EarlyStopping's wait state, or ReduceLROnPlateau's cooldown—need separate handling.
There is no documented public hook for adding arbitrary data to the BackupAndRestore payload, and its backup files should not be treated as an extension format. A wrapper can save selected public attributes and reapply them in on_train_begin, but that is a best-effort, non-atomic convention. Verify this against the installed TensorFlow release: record optimizer.iterations and the schedule value before interruption, then repeat after restore and compare callback counters.