How to save fine tuned PaliGemma model?

[SkalskiP]

Description of the feature request:

Hi 👋🏻

Thanks a lot for PaliGemma fine-tune example notebook. It worked without any issue. However, it lacks the last step - saving a fine-tuned model. (Or I'm just missing something.) I'd appreciate any help. Thank you. 🙏🏻

What problem are you trying to solve with this feature?

No response

Any other information you'd like to share?

No response

[gokayfem]

+1 We need this

[selamw1]

Hi @SkalskiP,

Based on this conversation at X, you can use numpy.savez, a function that efficiently stores multiple arrays in a single, uncompressed .npz format.

import big_vision.utils as bv_utils
flat, _ = bv_utils.tree_flatten_with_names(flat)
with open("ckpt.npz", "wb") as f:
  np.savez(f, **{k: v for k, v in flat})

And then load:

bv_utils.load_checkpoint_np("ckpt.npz")

[SkalskiP]

Thanks a lot! @gokayfem I really appreciate the help.

[gokayfem]

It should be

flat, _ = bv_utils.tree_flatten_with_names(params)