Flax: How to save a model and then reuse in Flax?

[varunsingh88]

Is there a way to save a trained model in Flax for using it later? I have found some solutions that aim at saving the model in a format that is suitable for import into tensorflow. But I am hoping to get a solution to save a Flax model and then re-use it in Flax.

[andsteing]

We usually use the module flax.training.checkpoints for this.

You can see usage in e.g. our examples/imagenet/train.py:

flax/examples/imagenet/train.py

Lines 197 to 206 in 02f53af

	def restore_checkpoint(state, workdir):
	return checkpoints.restore_checkpoint(workdir, state)
	def save_checkpoint(state, workdir):
	if jax.process_index() == 0:
	# get train state from the first replica
	state = jax.device_get(jax.tree_map(lambda x: x[0], state))
	step = int(state.step)
	checkpoints.save_checkpoint(workdir, state, step, keep=3)

[mtavit]

Hey, for some reason, I can't use flax.training.checkpoints in Kaggle. Do you happen to know why? I would really appreciate an answer!

[chiamp]

Hi @mtavit, do you have a snippet of code that can reproduce the error you're getting? What version of Flax are you using on Kaggle?

[mtavit]

Hey @chiamp , I am using version 0.6.4 of flax.

The code is just:
from flax.training import checkpoints

And this is the error:

---

ImportError Traceback (most recent call last)
/tmp/ipykernel_27/2471448347.py in
24 import flax
25 from flax import linen as nn
---> 26 from flax.training import checkpoints
27
28 # Optax

/opt/conda/lib/python3.7/site-packages/flax/training/checkpoints.py in
36 from flax import traverse_util
37 import jax
---> 38 from jax import monitoring
39 from jax import process_index
40 from jax import sharding

ImportError: cannot import name 'monitoring' from 'jax' (/opt/conda/lib/python3.7/site-packages/jax/init.py)

Which version of Flax do I need to work with checkpoints - and is it possible to get this in kaggle?

[chiamp]

I'm assuming you're using a Kaggle notebook. Under Notebook options, select the TPU VM v3-8 accelerator and switch the Internet option to be on.
Then run pip install jax==0.4.6 in the notebook to install the latest Jax version. Then execute from flax.training import checkpoints and there should be no error.

The issue was that the Kaggle runtime was using an older version of Jax. For some reason, the CPU (no accelerator) runtime only had 0.3.25 as the latest version of Jax, which isn't compatible with flax.training.checkpoints.

[mtavit]

Thank you very much!