tf.keras.callbacks.OrbaxCheckpoint

View source on GitHub

Callback to save and load model state using Orbax with a similar API to ModelCheckpoint.

Inherits From: Callback

tf.keras.callbacks.OrbaxCheckpoint(
    directory,
    monitor='val_loss',
    verbose=0,
    save_best_only=False,
    mode='auto',
    save_freq='epoch',
    initial_value_threshold=None,
    max_to_keep=1,
    save_on_background=True,
    save_weights_only=False
)

This callback saves the model's weights and optimizer state asynchronously using Orbax, allowing training to continue without blocking for I/O.

Multi-host Support: When running in a multi-host distributed training environment with JAX backend, this callback automatically coordinates checkpointing across all hosts to ensure consistency and proper synchronization. Multi-host checkpointing is only supported on JAX.

Example:

model.compile(loss=..., optimizer=..., metrics=['accuracy'])

EPOCHS = 10
checkpoint_dir = '/tmp/ckpt'
orbax_checkpoint_callback = keras.callbacks.OrbaxCheckpoint(
    directory=checkpoint_dir,
    monitor='val_accuracy',
    mode='max',
    save_best_only=True)

# Model is saved at the end of every epoch, if it's the best seen so far.
model.fit(epochs=EPOCHS, callbacks=[orbax_checkpoint_callback])

# Alternatively, save checkpoints every N batches -
orbax_checkpoint_callback = keras.callbacks.OrbaxCheckpoint(
    directory=checkpoint_dir,
    save_freq=100)  # Save every 100 batches

model.fit(epochs=EPOCHS, callbacks=[orbax_checkpoint_callback])

Methods

is_multihost_enabled

View source

is_multihost_enabled()

Return True if multi-host checkpointing is enabled and initialized.

This method can be used to check if the callback is operating in a multi-host distributed training environment. Multi-host checkpointing is only supported on JAX backend.

is_primary_host

View source

is_primary_host()

Return True if this process is the primary host in multi-host setup.

In multi-host environments, only the primary host typically handles logging and coordination tasks. Multi-host checkpointing is only supported on JAX backend.

on_batch_begin

View source

on_batch_begin(
    batch, logs=None
)

A backwards compatibility alias for on_train_batch_begin.

on_batch_end

View source

on_batch_end(
    batch, logs=None
)

A backwards compatibility alias for on_train_batch_end.

on_epoch_begin

View source

on_epoch_begin(
    epoch, logs=None
)

Called at the start of an epoch.

Subclasses should override for any actions to run. This function should only be called during TRAIN mode.

on_epoch_end

View source

on_epoch_end(
    epoch, logs=None
)

Called at the end of an epoch.

Subclasses should override for any actions to run. This function should only be called during TRAIN mode.

on_predict_batch_begin

View source

on_predict_batch_begin(
    batch, logs=None
)

Called at the beginning of a batch in predict methods.

Subclasses should override for any actions to run.

Note that if the steps_per_execution argument to compile in Model is set to N, this method will only be called every N batches.

on_predict_batch_end

View source

on_predict_batch_end(
    batch, logs=None
)

Called at the end of a batch in predict methods.

Subclasses should override for any actions to run.

Note that if the steps_per_execution argument to compile in Model is set to N, this method will only be called every N batches.

on_predict_begin

View source

on_predict_begin(
    logs=None
)

Called at the beginning of prediction.

Subclasses should override for any actions to run.

on_predict_end

View source

on_predict_end(
    logs=None
)

Called at the end of prediction.

Subclasses should override for any actions to run.

on_test_batch_begin

View source

on_test_batch_begin(
    batch, logs=None
)

Called at the beginning of a batch in evaluate methods.

Also called at the beginning of a validation batch in the fit methods, if validation data is provided.

Subclasses should override for any actions to run.

Note that if the steps_per_execution argument to compile in Model is set to N, this method will only be called every N batches.

on_test_batch_end

View source

on_test_batch_end(
    batch, logs=None
)

Called at the end of a batch in evaluate methods.

Also called at the end of a validation batch in the fit methods, if validation data is provided.

Subclasses should override for any actions to run.

Note that if the steps_per_execution argument to compile in Model is set to N, this method will only be called every N batches.

on_test_begin

View source

on_test_begin(
    logs=None
)

Called at the beginning of evaluation or validation.

Subclasses should override for any actions to run.

on_test_end

View source

on_test_end(
    logs=None
)

Called at the end of evaluation or validation.

Subclasses should override for any actions to run.

on_train_batch_begin

View source

on_train_batch_begin(
    batch, logs=None
)

Called at the beginning of a training batch in fit methods.

Subclasses should override for any actions to run.

Note that if the steps_per_execution argument to compile in Model is set to N, this method will only be called every N batches.

on_train_batch_end

View source

on_train_batch_end(
    batch, logs=None
)

Called at the end of a training batch in fit methods.

Subclasses should override for any actions to run.

Note that if the steps_per_execution argument to compile in Model is set to N, this method will only be called every N batches.

on_train_begin

View source

on_train_begin(
    logs=None
)

Called at the beginning of training.

Subclasses should override for any actions to run.

on_train_end

View source

on_train_end(
    logs=None
)

Called at the end of training.

Subclasses should override for any actions to run.

set_model

View source

set_model(
    model
)

set_params

View source

set_params(
    params
)

wait_until_finished

View source

wait_until_finished()

Wait for any in-progress checkpoint operations to complete. This method blocks until all asynchronous checkpoint save operations have completed across all hosts in a multi-host setup.