orbax_checkpoint_manager_save_restore_jax_arrays.py

1import os
2import shutil
3import jax
4import jax.numpy as jnp
5import orbax.checkpoint as ocp
6
7# 1. Prepare some dummy data to save
8state = {'a': jnp.arange(10), 'b': jnp.ones((5, 5))}
9path = '/tmp/orbax_quickstart/'
10
11# Clean up any existing directory from previous runs
12if os.path.exists(path):
13    shutil.rmtree(path)
14
15# 2. Initialize a CheckpointManager
16# CheckpointManager handles saving, loading, and managing multiple steps/checkpoints.
17options = ocp.CheckpointManagerOptions(max_to_keep=3, create=True)
18mngr = ocp.CheckpointManager(
19    path,
20    ocp.StandardCheckpointer(),  # Handles basic JAX/numpy types and PyTrees
21    options=options
22)
23
24# 3. Save the state
25# We provide the step number and the data to be saved.
26save_args = ocp.args.StandardSave(state)
27mngr.save(step=0, args=save_args)
28mngr.wait_until_finished()
29print(f"Checkpoint saved at step 0 to: {path}")
30
31# 4. Restore the state
32# To restore, we define the structure (target) we expect to load into.
33restore_args = ocp.args.StandardRestore(state)
34restored_state = mngr.restore(step=0, args=restore_args)
35
36print("Restored 'a':", restored_state['a'])
37print("Restored 'b':", restored_state['b'])