Is it possible to checkpoint a linen.Module instance?

[alvinsunyixiao]

It could be useful to save some of the hyperparameters stored in a linen.Module instance to a checkpoint, but it looks like linen.Module is a ordinary dataclass instead of a struct.dataclass. This means that orbax checkpoint API does not serialize a linen.Module instance. Here is a code snippet of failing to save the model to a checkpoint.

import jax
import jax.numpy as jnp
from flax import struct
from flax import linen as nn
from orbax.checkpoint import PyTreeCheckpointer

class Model(nn.Module):
    a: float  # some hyperparameter that I would want to store into a checkpoint

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


ckpter = PyTreeCheckpointer()
model = Model(a=10.0)
params = model.init(jax.random.PRNGKey(42), jnp.zeros(10))["params"]
ckpt = {
    "params": params,
    "model": model,
}
ckpter.save("/tmp/model", ckpt, force=True)
print(ckpter.restore("/tmp/model"))  # This gives {"params": ..., "model": None}

I tried the following changes

@struct.dataclass
class Model(nn.Module)

or

class Model(nn.Module, struct.PyTreeNode)

but neither of them worked. (I think the errors have something to do with Python dataclass inheritance problems)

Is there a way to make model also a struct.dataclass? And / Or is there a way to save the model to a checkpoint?

[cgarciae]

Hey @alvinsunyixiao, since Flax Modules are dataclasses you can use dataclasses.asdict to retrieve their config, e.g:

import flax.linen as nn
import dataclasses

module = nn.Dense(10)
module_config = dataclasses.asdict(module)

print(module_config)

Hope this helps!

[alvinsunyixiao]

Thank you very much! I will give that a try. Although is there reason why flax modules are not flax data classes?

[cgarciae]

struct.dataclass classes are pytrees, Modules are not for many good reasons.

[alvinsunyixiao]

But would it be possible to make a Module a pytree? So for example, being able to compute gradient with respect to some hyperparameters?

[cgarciae]

If you need something like this you could pass the hyper-parameters as input to the loss function being differentiated, construct the Module inside with those hyper-parameters, and call apply using the main params + inputs as usual. Alternatively, you could create new layers in which the hyperparameters of interest are actually params.

[alvinsunyixiao]

Thank you very much! I will give those suggestions a try!