-
Notifications
You must be signed in to change notification settings - Fork 305
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Orbax no longer supports lazy restoration of checkpoints.
If our checkpoint is `{'a': 2, 'b': 4}`, and we wish to 'lazily' load the tree such that only 'a' is materialized, equivalent functionality can be achieved using the following: ``` ckptr = ocp.PyTreeCheckpointer() tree = {'a': <dummy_value>, 'b': <dummy_value>} ckptr.restore(path, item=tree, transforms={'b': ocp.Transform(use_fallback=True)}) {'a': 2, 'b': <dummy_value>} ``` OR ``` ckptr = ocp.PyTreeCheckpointer() tree = {'a': <dummy_value>} # Returns a tree just containing 'a', materialized. ckptr.restore(path, item=tree, transforms={}) {'a': 2} ``` Previously this could have also been done with the lazy feature: ``` ckptr = ocp.PyTreeCheckpointer() ckptr.restore(path, restore_args={'a': RestoreArgs(), 'b': RestoreArgs(lazy=True)}) {'a': 2, 'b': LazyValue()} ``` PiperOrigin-RevId: 547941035
- Loading branch information
1 parent
eb08ffb
commit 831342f
Showing
5 changed files
with
36 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters