How to save/load model?

[dev-sora]

As the title says.
Of course I know there are low level API flax.serialization and some functions for checkpoint saving/restoring, but the document gives me only little information.
What I want to do is simply save model weights, and load it in another scripts. Can someone explain about that with practical example?

[matthias-wright]

Hi @dev-sora, here is an example:

import jax.numpy as jnp
import jax
import flax
import flax.linen as nn
from flax.training import train_state, checkpoints
import optax
import numpy as np


class Net(nn.Module):
    features: int

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(self.features)(x)
        x = nn.Dense(self.features)(x)
        return x


model = Net(features=2)
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 2)))

tx = optax.adam(learning_rate=0.0001)
state = train_state.TrainState.create(apply_fn=model.apply,
                                      params=params,
                                      tx=tx)


CKPT_DIR = 'ckpts'
checkpoints.save_checkpoint(ckpt_dir=CKPT_DIR, target=state, step=0)
restored_state = checkpoints.restore_checkpoint(ckpt_dir=CKPT_DIR, target=state)

assert jax.tree_util.tree_all(jax.tree_multimap(lambda x, y: (x == y).all(), state.params, restored_state.params))

[dev-sora]

@matthias-wright Thanks for your kind responding.
Now the only question is why you are targeting 'state' when restoring checkpoint. Does this mean checkpoints.restore_checkpoint receive the structure of model as 'target' argument, set the state of checkpoint, then return it as a new state object? If so, does it mean that I need to redefine and instantiate the class of the model and randomly initialize the state once I restore it in a different file than the one where I saved the checkpoint?

[matthias-wright]

You can also set target=None, but then restored_state will be a dict rather than a TrainState, see here.

[dev-sora]

Thanks!!

[dantp-ai]

The flax.training.checkpoints API is deprecated. Check flax's tutorial on how to use Orbax.