Skip to content

Internal change #212

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.

Sign up for GitHub

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

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+35 −3
Open

Internal change #212

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions recml/core/utils/keras_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -826,7 +826,7 @@ def _prepare_v3_restore(
# Load index path metadata for non-trainable and optimizer variables if
# needed.
non_trainable_paths = None
if model is not None and abstract_state.get(NON_TRAINABLE_VARIABLES_KEY):
if model is not None and NON_TRAINABLE_VARIABLES_KEY in abstract_state:
nt_path = epath.Path(checkpoint_path) / NON_TRAINABLE_PATHS_KEY
checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
try:
Expand All @@ -840,7 +840,7 @@ def _prepare_v3_restore(
if (
model is not None
and restore_optimizer_vars
and abstract_state.get(OPTIMIZER_VARIABLES_KEY)
and OPTIMIZER_VARIABLES_KEY in abstract_state
):
opt_path = epath.Path(checkpoint_path) / OPTIMIZER_PATHS_KEY
checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
Expand All @@ -859,7 +859,7 @@ def _prepare_v3_restore(
state_transforms = {}

for key in keys:
if not abstract_state.get(key):
if key not in abstract_state:
continue

filtered_abstract_state[key] = {}
Expand Down
32 changes: 32 additions & 0 deletions recml/core/utils/keras_utils_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,38 @@ def test_keras_orbax_checkpointer_v3(
# Ensures predictions are identical.
np.testing.assert_allclose(preds, preds_after_restoration)

def test_keras_orbax_checkpointer_v3_no_non_trainable_variables(self):
model = keras.Sequential(
[keras.layers.Dense(5, name="dense")], name="model"
)
model.build((2, 10))
model.compile(optimizer=keras.optimizers.Adam(0.01), loss="mse")
model.optimizer.build(model.trainable_variables)
self.assertEmpty(model.non_trainable_variables)

checkpoint_dir = self.create_tempdir().full_path
checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3(
checkpoint_dir, max_to_keep=5
)
checkpoint_manager.save_model_variables(model, 0)
checkpoint_manager.wait_until_finished()

new_model = keras.Sequential(
[keras.layers.Dense(5, name="dense")], name="model"
)
new_model.build((2, 10))
new_model.compile(optimizer=keras.optimizers.Adam(0.01), loss="mse")
new_model.optimizer.build(new_model.trainable_variables)

checkpoint_manager.restore_model_variables(new_model, 0)
checkpoint_manager.close()

for expected, observed in zip(
[v.value for v in model.trainable_variables],
[v.value for v in new_model.trainable_variables],
):
np.testing.assert_allclose(observed, expected)

def test_restore_keras_checkpoint(self):
dummy_inputs = _create_dummy_inputs()
bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs))
Expand Down
Loading