Need to preserve more than model weights? Orbax is a JAX-oriented library for checkpointing and serialization. It can store any JAX pytree, including arrays, dictionaries, metadata, configurations, and custom Flax dataclasses, while also supporting export to TensorFlow SavedModel. section 11.2.2
Its practical features include:
- Save and restore parameters with PyTreeCheckpointer. section 11.2.2
- Bundle smaller arrays with save_args to improve checkpointing performance. section 11.2.2
- Use CheckpointManager to save at intervals, limit retained checkpoints, and maintain checkpoint directories. section 11.2.2
- Support asynchronous checkpointing, multiple data types, and various storage formats. section 11.2.2
Saving parameters
Create a PyTreeCheckpointer, derive save arguments from the parameter pytree, save it to a filesystem path, and restore it later. A checkpoint manager can place the same operation inside a training loop, for example every ten steps while retaining four checkpoints. section 11.2.2
Listing 11.11 Saving and restoring model parameters
from flax.training import orbax_utils #1 import orbax.checkpoint #1 path = 'tmp/orbax/saved_model' #2 orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer() #3 save_args = orbax_utils.save_args_from_target(state.params) #4 orbax_checkpointer.save(path, state.params, save_args=save_args) #5 params_restored = orbax_checkpointer.restore(path) #6
The complete listing shows the minimal save-and-restore workflow with PyTreeCheckpointer and save_args. section 11.2.2
Listing 11.12 Using a checkpoint manager to regularly make checkpoints
orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer() #1
options = orbax.checkpoint.CheckpointManagerOptions( #2
max_to_keep=4, #3
save_interval_steps=10, #4
create=True) #5
checkpoint_manager = orbax.checkpoint.CheckpointManager( #6
'tmp/orbax/checkpoints', orbax_checkpointer, options)
# Inside a training loop
for epoch in range(100):
# ... do your training
checkpoint_manager.save(epoch, params, save_kwargs={'save_args':
save_args}) #7
The checkpoint-manager listing demonstrates interval-based saving and retention settings inside a training loop. section 11.2.2
