Internal change

recml/core/training/keras_trainer.py

@@ -90,6 +90,24 @@ def create_model_for_eval(self, **kwargs) -> keras.Model:
     """
     return self.create_model(**kwargs)
 
+  def create_eval_callbacks(
+      self, model: keras.Model
+  ) -> Sequence[keras.callbacks.Callback]:
+    """Returns extra callbacks for evaluating `model`.
+
+    They run after the trainer restores the evaluated checkpoint into `model`,
+    so `on_test_begin` sees its weights, e.g. to rebuild state derived from
+    them before the first evaluation step. Defaults to none.
+
+    Args:
+      model: The Keras model constructed by `create_model_for_eval`.
+
+    Returns:
+      A sequence of Keras callbacks.
+    """
+    del model
+    return []
+
   def export_model(self, model: keras.Model, model_dir: str):
     """Exports a Keras model.
 
@@ -375,6 +393,7 @@ def evaluate(self, task: KerasTask) -> core.Logs:
         **self._maybe_get_model_kws(task, dataset, val=True)
     )
 
+    callbacks = [*task.create_eval_callbacks(model), *self.eval_callbacks]
     if keras.backend.backend() == "jax":
       [tb_cbk] = [
           cbk
@@ -385,7 +404,7 @@ 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
@@ -398,7 +417,7 @@ def evaluate(self, task: KerasTask) -> core.Logs:
     return model.evaluate(
         dataset,
         steps=self._steps_per_eval,
-        callbacks=self.eval_callbacks,
+        callbacks=callbacks,
     )
 
   def train_and_evaluate(self, task: KerasTask) -> core.Logs:
@@ -488,6 +507,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None):
               f"Unsupported checkpoint version: {checkpoint_version}"
           )
 
+    task_callbacks = list(task.create_eval_callbacks(model))
     history = None
     for epoch in ocp.checkpoint_utils.checkpoints_iterator(
         self._checkpoint_dir,
@@ -500,7 +520,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, *self.eval_callbacks]
+      callbacks = [restore_callback, *task_callbacks, *self.eval_callbacks]
       try:
         logging.info(f"eval | epoch: {epoch: 6d} | {steps_msg}")
         epoch_start_time = time.time()

recml/core/training/keras_trainer_test.py

@@ -62,6 +62,25 @@ def on_test_end(self, logs=None):
     self._record(logs, "custom/marker", 7.0)
 
 
+class _RecordWeightsCallback(keras.callbacks.Callback):
+  """Records the model's kernel at the start of each evaluation."""
+
+  def __init__(self, model: keras.Model):
+    super().__init__()
+    self._model = model
+    self.kernels = []
+
+  def on_test_begin(self, logs=None):
+    self.kernels.append(float(self._model.layers[-1].kernel.numpy()[0, 0]))
+
+
+class _KerasTaskWithEvalCallbacks(_KerasTask):
+
+  def create_eval_callbacks(self, model: keras.Model):
+    self.callback = _RecordWeightsCallback(model)
+    return [self.callback]
+
+
 class KerasTrainerTest(parameterized.TestCase):
 
   def setUp(self):
@@ -219,6 +238,28 @@ def test_val_logs_warns_when_metrics_callback_overwrites_key(self):
     self.assertIn("overwrites existing evaluation log key", log.output[0])
     self.assertEqual(val_logs["val_custom/marker"], 7.0)
 
+  def test_task_eval_callbacks_see_restored_weights(self):
+    if keras.backend.backend() != "jax":
+      self.skipTest("Continuous eval is only supported on the Jax backend.")
+    trainer = keras_trainer.KerasTrainer(
+        distribution=keras.distribution.DataParallel(),
+        train_steps=5,
+        steps_per_eval=1,
+        steps_per_loop=5,
+        model_dir=self.create_tempdir().full_path,
+        continuous_eval_timeout=5,
+    )
+    task = _KerasTaskWithEvalCallbacks()
+    experiment = core.Experiment(task, trainer)
+    core.run_experiment(experiment, core.Trainer.Mode.TRAIN)
+    keras.backend.clear_session()
+
+    core.run_experiment(experiment, core.Trainer.Mode.CONTINUOUS_EVAL)
+    # The kernel starts at -1.0 and training moves it, so the callback ran
+    # after the checkpoint restore.
+    self.assertNotEmpty(task.callback.kernels)
+    self.assertNotAlmostEqual(task.callback.kernels[0], -1.0)
+
   def test_eval_name(self):
     if keras.backend.backend() != "jax":
       self.skipTest(