Skip to content

Internal change #197

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 #197

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
78 changes: 66 additions & 12 deletions recml/core/training/keras_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from __future__ import annotations

import abc
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
import dataclasses
import functools
import gc
Expand Down Expand Up @@ -102,6 +102,25 @@ def export_model(self, model: keras.Model, model_dir: str):
model_dir: The model directory passed to the trainer.
"""

def create_callbacks(
self, training: bool
) -> Sequence[keras.callbacks.Callback]:
"""Creates task-specific callbacks for training or evaluation.

The trainer runs these callbacks before its own callbacks. Thus metrics that
they add to `logs` are written to TensorBoard. In evaluation, use
`keras_utils.MetricsCallback` to report metrics, because Keras does not
return changes made to `logs` in `on_test_end`.

Args:
training: Whether the callbacks are for training or for evaluation.

Returns:
The callbacks. The default is no callbacks.
"""
del training
return ()


class KerasTrainer(core.Trainer[KerasTask]):
"""Keras 3 trainer for use with the Keras task."""
Expand Down Expand Up @@ -307,6 +326,29 @@ def _maybe_get_model_kws(

return kws

def _task_callbacks(
self, task: KerasTask, *, training: bool
) -> list[keras.callbacks.Callback]:
"""Returns the task callbacks, or none if the task has no hook."""
create_callbacks = getattr(task, "create_callbacks", None)
if create_callbacks is None:
return []
return list(create_callbacks(training=training))

def _val_logs(
self,
history: Mapping[str, Any],
callbacks: Sequence[keras.callbacks.Callback],
steps_per_second: float,
) -> dict[str, Any]:
"""Returns the logs for `EpochSummaryCallback` after an evaluation."""
val_logs = {"val_" + k: v for k, v in history.items()}
for cbk in callbacks:
if isinstance(cbk, keras_utils.MetricsCallback):
val_logs.update({"val_" + k: v for k, v in cbk.latest_metrics.items()})
val_logs["val_steps_per_second"] = steps_per_second
return val_logs

def train(self, task: KerasTask) -> core.Logs:
"""Trains a Keras model."""
dataset = task.create_dataset(
Expand All @@ -320,7 +362,10 @@ def train(self, task: KerasTask) -> core.Logs:
dataset,
epochs=self._train_epochs + 1,
steps_per_epoch=self._steps_per_loop,
callbacks=self.train_callbacks,
callbacks=[
*self._task_callbacks(task, training=True),
*self.train_callbacks,
],
initial_epoch=1,
)
model.summary(print_fn=logging.info)
Expand All @@ -339,6 +384,10 @@ def evaluate(self, task: KerasTask) -> core.Logs:
**self._maybe_get_model_kws(task, dataset, val=True)
)

callbacks = [
*self._task_callbacks(task, training=False),
*self.eval_callbacks,
]
if keras.backend.backend() == "jax":
[tb_cbk] = [
cbk
Expand All @@ -349,20 +398,20 @@ def evaluate(self, task: KerasTask) -> core.Logs:
history = model.evaluate(
dataset,
steps=self._steps_per_eval,
callbacks=self.eval_callbacks,
callbacks=callbacks,
return_dict=True,
)
epoch_dt = time.time() - epoch_start_time
steps_per_second = self._steps_per_eval / epoch_dt # pyrefly: ignore[unsupported-operation]
val_logs = {"val_" + k: v for k, v in history.items()}
val_logs["val_steps_per_second"] = steps_per_second
tb_cbk.on_epoch_end(0, val_logs)
tb_cbk.on_epoch_end(
0, self._val_logs(history, callbacks, steps_per_second)
)
return history

return model.evaluate(
dataset,
steps=self._steps_per_eval,
callbacks=self.eval_callbacks,
callbacks=callbacks,
)

def train_and_evaluate(self, task: KerasTask) -> core.Logs:
Expand All @@ -388,7 +437,10 @@ def train_and_evaluate(self, task: KerasTask) -> core.Logs:
steps_per_epoch=self._steps_per_loop,
# Explicitly set to None for deterministic evaluation.
validation_steps=None,
callbacks=self.train_callbacks,
callbacks=[
*self._task_callbacks(task, training=True),
*self.train_callbacks,
],
initial_epoch=1,
)
model.summary(print_fn=logging.info)
Expand Down Expand Up @@ -453,6 +505,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
)

history = None
task_callbacks = self._task_callbacks(task, training=False)
for epoch in ocp.checkpoint_utils.checkpoints_iterator(
self._checkpoint_dir,
timeout=self._continuous_eval_timeout,
Expand All @@ -464,6 +517,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
for cbk in self.eval_callbacks
if isinstance(cbk, keras_utils.EpochSummaryCallback)
]
callbacks = [restore_callback, *task_callbacks, *self.eval_callbacks]
try:
logging.info(f"eval | epoch: {epoch: 6d} | {steps_msg}")
epoch_start_time = time.time()
Expand All @@ -475,7 +529,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
history = model.evaluate(
eval_dataset,
steps=self._steps_per_eval,
callbacks=[restore_callback] + self.eval_callbacks,
callbacks=callbacks,
return_dict=True,
)

Expand All @@ -491,9 +545,9 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
epoch_dt = time.time() - epoch_start_time
steps_per_second = self._steps_per_eval / epoch_dt # pyrefly: ignore[unsupported-operation]

val_logs = {"val_" + k: v for k, v in history.items()}
val_logs["val_steps_per_second"] = steps_per_second
tb_cbk.on_epoch_end(epoch, val_logs)
tb_cbk.on_epoch_end(
epoch, self._val_logs(history, callbacks, steps_per_second)
)
except FileNotFoundError:
logging.info("Checkpoint epoch: %s did not finish writing...", epoch)

Expand Down
82 changes: 82 additions & 0 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 @@ -52,6 +52,27 @@ def create_model(self) -> keras.Model:
return model


class _MarkerCallback(keras_utils.MetricsCallback):
"""Reports a fixed metric at the end of each training epoch and eval pass."""

def on_epoch_end(self, epoch, logs=None):
self._record(logs, "custom/marker", 7.0)

def on_test_end(self, logs=None):
self._record(logs, "custom/marker", 7.0)


class _KerasTaskWithCallbacks(_KerasTask):

def __init__(self):
self.created_callbacks = []

def create_callbacks(self, training: bool):
callback = _MarkerCallback()
self.created_callbacks.append((training, callback))
return [callback]


class KerasTrainerTest(parameterized.TestCase):

def setUp(self):
Expand Down Expand Up @@ -129,6 +150,67 @@ def test_keras_task_and_trainer(
):
self.assertEqual(history.history["num_params/trainable"][0], 2)

def test_task_callbacks_run_in_training(self):
trainer = keras_trainer.KerasTrainer(
distribution=(
keras.distribution.DataParallel()
if keras.backend.backend() == "jax"
else None
),
train_steps=4,
steps_per_loop=2,
model_dir=self.create_tempdir().full_path,
)
task = _KerasTaskWithCallbacks()

history = core.run_experiment(
core.Experiment(task, trainer), core.Trainer.Mode.TRAIN
)

self.assertEqual([t for t, _ in task.created_callbacks], [True])
self.assertEqual(history.history["custom/marker"], [7.0, 7.0])

@parameterized.named_parameters(
{"testcase_name": "eval", "mode": core.Trainer.Mode.EVAL},
{
"testcase_name": "continuous_eval",
"mode": core.Trainer.Mode.CONTINUOUS_EVAL,
},
)
def test_task_metrics_callbacks_are_written_in_eval(self, mode):
if keras.backend.backend() != "jax":
self.skipTest("`EpochSummaryCallback` is only used on the Jax backend.")
model_dir = self.create_tempdir().full_path

def _trainer():
return keras_trainer.KerasTrainer(
distribution=keras.distribution.DataParallel(),
train_steps=2,
steps_per_eval=1,
steps_per_loop=2,
model_dir=model_dir,
continuous_eval_timeout=1,
)

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

with mock.patch.object(
keras_utils.EpochSummaryCallback, "on_epoch_end", autospec=True
) as on_epoch_end:
core.run_experiment(core.Experiment(task, _trainer()), mode)

self.assertEqual([t for t, _ in task.created_callbacks], [False])
val_logs = on_epoch_end.call_args.args[2]
self.assertEqual(val_logs["val_custom/marker"], 7.0)
self.assertIn("val_loss", val_logs)
self.assertIn("val_steps_per_second", val_logs)

def test_eval_name(self):
if keras.backend.backend() != "jax":
self.skipTest(
Expand Down
Loading
Loading