callback to back up and restore the training state backupandrestore callback is intended to recover training from an interruption that has happened in the middle of a model fit execution by backing up the training states in a temporary checkpoint file at the end of each epoch each backup overwrites the previously written checkpoint file so at any given time there is at most one such checkpoint file for backuprestoring purpose if training restarts before completion the training state which includes the model weights and epoch number is restored to the most recently saved state at the beginning of a new model fit run at the completion of a model fit run the temporary checkpoint file is deleted note that the user is responsible to bring jobs back after the interruption this callback is important for the backup and restore mechanism for fault tolerance purpose and the model to be restored from a previous checkpoint is expected to be the same as the one used to back up if user changes arguments passed to compile or fit the checkpoint saved for fault tolerance can become invalid example class interruptingcallbackkeras callbacks callback def onepochbeginself epoch logsnone if epoch 4 raise runtimeerror interrupting callback keras callbacks backupandrestorebackupdirtmpbackup model keras models sequentialkeras layers dense10 model compilekeras optimizers sgd loss mse try model fitnp arange100 reshape5 20 np zeros5 epochs10 batchsize1 callbackscallback interruptingcallback verbose0 except pass history model fitnp arange100 reshape5 20 np zeros5 epochs10 batchsize1 callbackscallback verbose0 only 6 more epochs are run since first training got interrupted at zeroindexed epoch 4 second training will continue from 4 to 9 lenhistory history loss 6 args backupdir string path of directory where to store the data needed to restore the model the directory cannot be reused elsewhere to store other files e g by the backupandrestore callback of another training run or by another callback e g modelcheckpoint of the same training run savefreq epoch integer or false when set to epoch the callback saves the checkpoint at the end of each epoch when set to an integer the callback saves the checkpoint every savefreq batches set savefreqfalse only if using preemption checkpointing i e with savebeforepreemptiontrue deletecheckpoint boolean defaults to true this backupandrestore callback works by saving a checkpoint to back up the training state if deletecheckpointtrue the checkpoint will be deleted after training is finished use false if you d like to keep the checkpoint for future usage get training state from temporary file and restore it if fileutils existsself weightspath self model loadweightsself weightspath def onepochendself epoch logsnone self currentepoch epoch if self savefreq epoch self savemodel def ontrainbatchendself batch logsnone if self shouldsaveonbatchbatch self savemodel def savemodelself create host directory if it doesn t exist if not fileutils existsself backupdir fileutils makedirsself backupdir self model saveweightsfilepathself weightspath overwritetrue def shouldsaveonbatchself batch callback to back up and restore the training state backupandrestore callback is intended to recover training from an interruption that has happened in the middle of a model fit execution by backing up the training states in a temporary checkpoint file at the end of each epoch each backup overwrites the previously written checkpoint file so at any given time there is at most one such checkpoint file for backup restoring purpose if training restarts before completion the training state which includes the model weights and epoch number is restored to the most recently saved state at the beginning of a new model fit run at the completion of a model fit run the temporary checkpoint file is deleted note that the user is responsible to bring jobs back after the interruption this callback is important for the backup and restore mechanism for fault tolerance purpose and the model to be restored from a previous checkpoint is expected to be the same as the one used to back up if user changes arguments passed to compile or fit the checkpoint saved for fault tolerance can become invalid example class interruptingcallback keras callbacks callback def on_epoch_begin self epoch logs none if epoch 4 raise runtimeerror interrupting callback keras callbacks backupandrestore backup_dir tmp backup model keras models sequential keras layers dense 10 model compile keras optimizers sgd loss mse try model fit np arange 100 reshape 5 20 np zeros 5 epochs 10 batch_size 1 callbacks callback interruptingcallback verbose 0 except pass history model fit np arange 100 reshape 5 20 np zeros 5 epochs 10 batch_size 1 callbacks callback verbose 0 only 6 more epochs are run since first training got interrupted at zero indexed epoch 4 second training will continue from 4 to 9 len history history loss 6 args backup_dir string path of directory where to store the data needed to restore the model the directory cannot be reused elsewhere to store other files e g by the backupandrestore callback of another training run or by another callback e g modelcheckpoint of the same training run save_freq epoch integer or false when set to epoch the callback saves the checkpoint at the end of each epoch when set to an integer the callback saves the checkpoint every save_freq batches set save_freq false only if using preemption checkpointing i e with save_before_preemption true delete_checkpoint boolean defaults to true this backupandrestore callback works by saving a checkpoint to back up the training state if delete_checkpoint true the checkpoint will be deleted after training is finished use false if you d like to keep the checkpoint for future usage get training state from temporary file and restore it saves the model args epoch the epoch this iteration is in batch the batch this iteration is in none if the save_freq is set to epoch logs the logs dict passed in to on_batch_end or on_epoch_end create host directory if it doesn t exist handles batch level saving logic supports steps_per_execution new epoch batches are zero indexed
from keras.api_export import keras_export
from keras.callbacks.callback import Callback
from keras.utils import file_utils


@keras_export("keras.callbacks.BackupAndRestore")
class BackupAndRestore(Callback):
    

    def __init__(
        self,
        backup_dir,
        save_freq="epoch",
        delete_checkpoint=True,
    ):
        super().__init__()
        self.save_freq = save_freq
        self.delete_checkpoint = delete_checkpoint
        self._batches_seen_since_last_saving = 0
        self._last_batch_seen = 0
        self._current_epoch = 0

        if not backup_dir:
            raise ValueError("Empty `backup_dir` argument passed")
        self.backup_dir = backup_dir
        self._weights_path = file_utils.join(backup_dir, "latest.weights.h5")
        if save_freq != "epoch" and not isinstance(save_freq, int):
            raise ValueError(
                "Invalid value for argument `save_freq`. "
                f"Received: save_freq={save_freq}. "
                "Expected either 'epoch' or an integer value."
            )

    def on_train_begin(self, logs=None):
        
        if file_utils.exists(self._weights_path):
            self.model.load_weights(self._weights_path)

    def on_epoch_end(self, epoch, logs=None):
        self._current_epoch = epoch
        if self.save_freq == "epoch":
            self._save_model()

    def on_train_batch_end(self, batch, logs=None):
        if self._should_save_on_batch(batch):
            self._save_model()

    def _save_model(self):
        
        
        if not file_utils.exists(self.backup_dir):
            file_utils.makedirs(self.backup_dir)
        self.model.save_weights(filepath=self._weights_path, overwrite=True)

    def _should_save_on_batch(self, batch):
        
        if self.save_freq == "epoch":
            return False
        if batch <= self._last_batch_seen:  
            add_batches = batch + 1  
        else:
            add_batches = batch - self._last_batch_seen
        self._batches_seen_since_last_saving += add_batches
        self._last_batch_seen = batch

        if self._batches_seen_since_last_saving >= self.save_freq:
            self._batches_seen_since_last_saving = 0
            return True
        return False

    def on_train_end(self, logs=None):
        if self.delete_checkpoint and file_utils.exists(self.backup_dir):
            file_utils.rmtree(self.backup_dir)