Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

JAX checkpointing via Orbax #1491

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

Conversation

tttc3
Copy link

@tttc3 tttc3 commented Sep 25, 2023

Split from #1490 . Provides a very basic checkpointing functionality for flax based JAX models via Orbax. Essentially follows the simple example from flax documentation, without the use of the context manager.

As noted in #1490 , Orbax is currently imported lazily (only imported once first called) in a non PEP8 compliant manner, to avoid requiring Orbax as a package dependency. I feel that there should be a better way to do this, but I can't currently think of one.
As such, I've marked this as a draft for now.

Any thought or suggestions are greatly appreciated!

@lululxvi
Copy link
Owner

From the information here https://orbax.readthedocs.io/en/latest/, Orbax is not a package, but rather a separate package for each functionality provided by the Orbax namespace. For checkpoint, we need orbax-checkpoint

@tttc3
Copy link
Author

tttc3 commented Sep 25, 2023

I've updated the import to only require orbax.checkpointing

@lululxvi
Copy link
Owner

As checkpoint is a very basic tool, I think it is ok to add it to the required package dependency.

@tttc3
Copy link
Author

tttc3 commented Sep 25, 2023

Added orbax-checkpoint as a top level import and a package wide dependency. Still have flax.training as a non top level import though.

@@ -1,5 +1,6 @@
matplotlib
numpy
orbax-checkpoint
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need to move this to the JAX part below.

@@ -37,6 +37,7 @@ classifiers = [
dependencies = [
"matplotlib",
"numpy",
"orbax-checkpoint",
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not add here.

@@ -1,5 +1,6 @@
matplotlib
numpy
orbax-checkpoint
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not add here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants