Saving an optimizer?

[HHalva]

Is there a recommended way for storing and loading an optimizer?

[cgarciae]

cloudpickle works fine.

[HHalva]

Yes I did that but I meant can I somehow extract the epoch number from opt_state's inner counter so as to explicitly get the epoch number. For example something like this:

opt_state = pickle.load(saved_opt_state)
iteration_resume = ? get this some how from opt_state 
for it in range(iteration_resume, num_iterations):
    params, opt_state = training_step(params, opt_state, it)

point being I need to also explicitly pass it rather than just having a counter within opt_state so would be useful if I can somehow get it from the opt_state.

[cgarciae]

Maybe just bundle stuff in another object and pickle that object with everything?

[cgarciae]

You could even just pickle/unpickle a (epoch, opt_state, optimizer) tuple.

[HHalva]

Yeah I did something like that in the end, was just curious if I could have gotten it directly from opt_state, but this works fine too -- thanks!

[rosshemsley]

For some optimizers, a count of the number of steps is maintained in the state itself. For example, adam keeps "count" in the ScaleByAdamState:

opt = optax.adamw(learning_rate=1e-4)
opt_state = opt.init(params)

# prints 0
print(opt_state[0].count)
updates, opt_state = opt.update(params, opt_state, params)

# prints 1
print(opt_state[0].count)
updates, opt_state = opt.update(params, opt_state, params)

# prints 2
print(opt_state[0].count)
updates, opt_state = opt.update(params, opt_state, params)

However, this count is not guaranteed to be present in every optimizer, and so we recommend checkpointing the data loader state alongside the optimizer state as @cgarciae suggested above!

[guyujun]

load params you need:

def load_opt_state_from_pkl(pkl_path):
    params = pickle.load(open(pkl_path, "rb"))
    _opt_state = params["state"]      #jopt.pack_optimizer_state(params)
    lr = params["lr"]
    # create opt_state obj.
    opt_state = [optax.ScaleByAdamState(count=jnp.array(1), mu={}, nu={}), 
                optax.ScaleByScheduleState(count=jnp.array(1))]
    opt_state[0].mu.update(_opt_state[0].mu)
    opt_state[0].nu.update(_opt_state[0].nu)
    opt_state = tuple(opt_state)
    return opt_state, lr

[aaarrti]

I would typically save my complete flax.training.TrainState using orbax-checkpoint. Then, before training

from orbax.checkpoint import CheckpointManager, checkpoint_utils
from jax import tree_util
import optax

TX = TypeVar("TX", bound=optax.OptState)


def restore_optimizer_state(opt_state: TX, restored: Mapping[str, ...]) -> TX:
    """Restore optimizer state from loaded checkpoint (or .msgpack file)."""
    return tree_util.tree_unflatten(
        tree_util.tree_structure(opt_state), tree_util.tree_leaves(restored)
    )

# Make "empty" training state
training_state = ...

mngr = CheckpointManager(...)
# Write checkpoint during previous training, which was interrupted

all_steps = mngr.all_steps(True)
latest_step = max(all_steps)
restore_args = checkpoint_utils.construct_restore_args(training_state)
restored_dict = mngr.restore(latest_step, restore_kwargs={"restore_args": restore_args})  
restored_optimizer = restore_optimizer_state(training_state.opt_state, restored_dict["opt_state"])
training_state.replace(params=restored_dict["params"], step=restored_dict["step"], opt_state=restored_optimizer, ...)  

from flax.serialization import from_bytes, msgpack_restore, to_bytes

If you would rather not use orbax, the same approach also works with flax.serialization and msgpack

from flax.serialization import msgpack_restore, to_bytes
import zlib

def save_as_msgpack(params, save_path: str, compression = None) -> None:
    msgpack_bytes: bytes = to_bytes(params)
    if compression == "GZIP":
        msgpack_bytes = zlib.compress(msgpack_bytes)
    with open(save_path, "wb+") as file:
        file.write(msgpack_bytes)


def load_from_msgpack(params, save_path: str, compression = None) -> Dict[str, Any]:
    bytes_data = file.read()
    if compression == "GZIP":
        bytes_data = zlib.decompress(bytes_data)

    params = msgpack_restore(bytes_data)
    return params

# You can dump/reload any pytree using this functions. 
# You'd still need to create an "empy" optimizer, to restore actuall pytree, 
# since reloaded object will be a plain python dict.