Can I save/restore part of the state?

[davidshen84]

Hi,

For models like MAE or any GAN-like models, parts of the model are only used for training but not during inferring. I wonder if I can save/restore only part of the state using the checkpoint manager.

E.g. my MAE trainer model comprises several smaller nn.Module instances.

mae_trainer = MAE(Encoder(...), Decoder(..), ...)

Later, when I use the Encoder instance, I want to from the checkpoint file only the part that relates to the Encoder model.

[peterdavidfagan]

Hi @davidshen84,

It looks like with Orbax you can save and restore PyTrees directly, this is outlined in the following doc. As far as my understanding goes you could use the outlined methods to save the parameters you care about using an instance of orbax.checkpoint.PyTreeCheckpointer().

The following block of code from the docs seems to provide and example:

checkpointer = orbax.checkpoint.PyTreeCheckpointer()
# 'path/to/directory' should already exist, but 'checkpoint_name' folder should
# not.
# The leaves of `my_tree` may be a number of different types.
# See `PyTreeCheckpointHandler` documentation.
checkpointer.save('path/to/directory/checkpoint_name/', my_tree)
# If you want to restore any of the leaves as sharded arrays, you'll need some
# extra arguments. See `PyTreeCheckpointHandler` documentation.
restored = checkpointer.restore('path/to/directory/checkpoint_name/')

[peterdavidfagan]

I am also looking into PyTree. But I could not figure out how to do it. Do you have an example?

It should be possible to access your model parameters, these should be a dict with the latest version of the library, this dict is technically a PyTree as per the definition (based on my understanding).

It should be possible to subset the returned parameters for those you wish to save. Once you have this subset saving should follow the API calls from orbax given above.

I'll write a colab if I get the time later. The above is my current understanding but I honestly need to verify with an example.

[davidshen84]

I think partial saving is easy.

checkpoint={
  'state': train_state,
  'encoder': train_state.params['encoder']
}

Then, the next problem is how to load this encoder parameter into the Encoder model. I am reading the discussion at #1619. But I have not worked it out yet.

Also, the encoder's parameters should not be updated when training in the downstream model. This thread, #1453, is also talking about this.

In summary, the checkpoint manager provided easy to use API to save and restore the train state. But we don't have an easy API to save and restore a model...really?

[peterdavidfagan]

Hey @davidshen84,

Apologies I have been quite busy since this post. I might get a chance later this week to look into the too. Thank you for sharing the above threads. I think a toy example in a Google Colab would be nice. I'd be happy to start one and link it here but I don't think I'll get to contribute to this until later this week.

[davidshen84]

Hi @peterdavidfagan ,

No worries. This document helped me a lot. I got some of my code snippet working.

[peterdavidfagan]

That document is very helpful.

[davidshen84]

Of course 😊

…

On Thu, 10 Aug 2023, 18:18 Peter David Fagan, ***@***.***> wrote: Hi @davidshen84 <https://github.com/davidshen84>, If I write a colab would you be able to accept the answer, this will help me to get the galaxy brain badge on GitHub. — Reply to this email directly, view it on GitHub <#3263 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAAQBTJU6HZIPVDX4KOBAD3XUSKMTANCNFSM6AAAAAA3GSQEQU> . You are receiving this because you were mentioned.Message ID: ***@***.***>