Skip to content

Internal change #168

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
Open

Internal change #168

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
79 changes: 63 additions & 16 deletions recml/core/training/keras_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,9 +122,32 @@ def __init__(
max_checkpoints_to_keep: int = 5,
checkpoint_save_interval_epochs: int = 1,
rng_seed: int = core.DEFAULT_RNG_SEED,
legacy_checkpoint_format: bool = True,
checkpoint_version: keras_utils.CheckpointVersion | str = "v2",
legacy_checkpoint_format: bool | None = None,
):
"""Initializes the instance."""
"""Initializes the instance.

Args:
distribution: The distribution strategy to use.
model_dir: The directory to save checkpoints and logs.
train_steps: Total number of training steps.
steps_per_eval: Number of steps between evaluations.
continuous_eval_timeout: Timeout for continuous evaluation.
steps_per_loop: Number of steps per training loop.
max_checkpoints_to_keep: Maximum number of checkpoints to keep.
checkpoint_save_interval_epochs: Interval in epochs to save checkpoints.
rng_seed: Random seed.
checkpoint_version: The checkpoint version to use. Supported versions:
"v1" (legacy V1), "v2" (V2, default), "v3" (V3).
legacy_checkpoint_format: Deprecated. Use checkpoint_version instead. If
set, True maps to V1, and False maps to V2. TODO(b/542602169): Remove
this in v2.
resume_training_launcher: Launcher to resume training.
enable_xmanager_measurements: Whether to enable XManager measurements.
enable_autoxprof: Whether to enable AutoXprof.
autoxprof_settings: Settings for AutoXprof.
s2_logging_settings: Settings for S2 logging.
"""

keras.utils.set_random_seed(rng_seed)

Expand Down Expand Up @@ -153,25 +176,42 @@ def __init__(
self._checkpoint_dir = os.path.join(model_dir, core.CHECKPOINT_DIR)
self._max_checkpoints_to_keep = max_checkpoints_to_keep
self._checkpoint_save_interval_epochs = checkpoint_save_interval_epochs
self._legacy_checkpoint_format = legacy_checkpoint_format
if legacy_checkpoint_format is not None:
logging.warning(
"legacy_checkpoint_format is deprecated, use checkpoint_version"
" instead."
)
self._checkpoint_version = (
keras_utils.CheckpointVersion.V1
if legacy_checkpoint_format
else keras_utils.CheckpointVersion.V2
)
else:
self._checkpoint_version = keras_utils.CheckpointVersion(
checkpoint_version
)

@functools.cached_property
def train_callbacks(self) -> list[keras.callbacks.Callback]:
"""Returns the training callbacks."""
if keras.backend.backend() == "jax":
if self._legacy_checkpoint_format:
checkpoint_manager = keras_utils.KerasOrbaxCheckpointManager(
checkpoint_dir=self._checkpoint_dir,
max_to_keep=self._max_checkpoints_to_keep,
save_interval_epochs=self._checkpoint_save_interval_epochs,
)
if self._checkpoint_version == keras_utils.CheckpointVersion.V1:
manager_cls = keras_utils.KerasOrbaxCheckpointManager
elif self._checkpoint_version == keras_utils.CheckpointVersion.V2:
manager_cls = keras_utils.KerasOrbaxCheckpointManagerV2
elif self._checkpoint_version == keras_utils.CheckpointVersion.V3:
manager_cls = keras_utils.KerasOrbaxCheckpointManagerV3
else:
checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV2(
checkpoint_dir=self._checkpoint_dir,
max_to_keep=self._max_checkpoints_to_keep,
save_interval_epochs=self._checkpoint_save_interval_epochs,
raise ValueError(
f"Unsupported checkpoint version: {self._checkpoint_version}"
)

checkpoint_manager = manager_cls(
checkpoint_dir=self._checkpoint_dir,
max_to_keep=self._max_checkpoints_to_keep,
save_interval_epochs=self._checkpoint_save_interval_epochs,
)

callbacks = [
keras_utils.EpochSummaryCallback(
log_dir=os.path.join(self._model_dir, core.LOG_DIR),
Expand Down Expand Up @@ -379,7 +419,7 @@ def timeout_fn() -> bool:
else:
steps_msg = "running complete evaluation..."

use_legacy_checkpoint_format = self._legacy_checkpoint_format
checkpoint_version = self._checkpoint_version

class _RestoreCallback(keras.callbacks.Callback):
"""Callback for restoring the model from the latest checkpoint."""
Expand All @@ -393,14 +433,21 @@ def __init__(
self._epoch = epoch

def on_test_begin(self, logs: Mapping[str, Any] | None = None):
if use_legacy_checkpoint_format:
if checkpoint_version == keras_utils.CheckpointVersion.V1:
keras_utils.restore_keras_model(
model, self._checkpoint_dir, step=self._epoch
)
else:
elif checkpoint_version in (
keras_utils.CheckpointVersion.V2,
keras_utils.CheckpointVersion.V3,
):
keras_utils.restore_keras_checkpoint(
self._checkpoint_dir, model=model, epoch=self._epoch
)
else:
raise ValueError(
f"Unsupported checkpoint version: {checkpoint_version}"
)

history = None
for epoch in ocp.checkpoint_utils.checkpoints_iterator(
Expand Down
32 changes: 23 additions & 9 deletions recml/core/training/keras_trainer_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ class KerasTrainerTest(parameterized.TestCase):

def setUp(self):
super().setUp()
keras.backend.clear_session()
# Workaround to make `create_tempdir` work with pytest.
if not flags.FLAGS.is_parsed():
flags.FLAGS.mark_as_parsed()
Expand All @@ -63,26 +64,38 @@ def setUp(self):
{"testcase_name": "train", "mode": core.Trainer.Mode.TRAIN},
{"testcase_name": "eval", "mode": core.Trainer.Mode.EVAL},
{
"testcase_name": "train_and_eval",
"testcase_name": "train_and_eval_v3",
"mode": core.Trainer.Mode.TRAIN_AND_EVAL,
"checkpoint_version": "v3",
},
{
"testcase_name": "continuous_eval_",
"mode": core.Trainer.Mode.CONTINUOUS_EVAL,
"testcase_name": "train_and_eval_v2",
"mode": core.Trainer.Mode.TRAIN_AND_EVAL,
"checkpoint_version": "v2",
},
{
"testcase_name": "train_and_eval_legacy_checkpoint_format",
"testcase_name": "train_and_eval_v1",
"mode": core.Trainer.Mode.TRAIN_AND_EVAL,
"legacy_checkpoint_format": True,
"checkpoint_version": "v1",
},
{
"testcase_name": "continuous_eval_v3",
"mode": core.Trainer.Mode.CONTINUOUS_EVAL,
"checkpoint_version": "v3",
},
{
"testcase_name": "continuous_eval_v2",
"mode": core.Trainer.Mode.CONTINUOUS_EVAL,
"checkpoint_version": "v2",
},
{
"testcase_name": "continuous_eval_legacy_checkpoint_format",
"testcase_name": "continuous_eval_v1",
"mode": core.Trainer.Mode.CONTINUOUS_EVAL,
"legacy_checkpoint_format": True,
"checkpoint_version": "v1",
},
)
def test_keras_task_and_trainer(
self, mode: str, legacy_checkpoint_format: bool = False
self, mode: str, checkpoint_version: str = "v3"
):
if keras.backend.backend() == "jax":
distribution = keras.distribution.DataParallel()
Expand All @@ -98,13 +111,14 @@ def test_keras_task_and_trainer(
steps_per_loop=2,
model_dir=self.create_tempdir().full_path,
continuous_eval_timeout=5,
legacy_checkpoint_format=legacy_checkpoint_format,
checkpoint_version=checkpoint_version,
)
experiment = core.Experiment(_KerasTask(), trainer)

if mode == core.Trainer.Mode.CONTINUOUS_EVAL:
# Produce one checkpoint so there is something to evaluate.
core.run_experiment(experiment, core.Trainer.Mode.TRAIN)
keras.backend.clear_session()

history = core.run_experiment(experiment, mode)

Expand Down
Loading
Loading