diff --git a/recml/core/training/keras_trainer.py b/recml/core/training/keras_trainer.py index dcb5736..bc9adbb 100644 --- a/recml/core/training/keras_trainer.py +++ b/recml/core/training/keras_trainer.py @@ -126,6 +126,8 @@ def __init__( legacy_checkpoint_format: bool | None = None, train_callbacks: Sequence[keras.callbacks.Callback] = (), eval_callbacks: Sequence[keras.callbacks.Callback] = (), + verbose: int | str = "auto", + steps_per_execution: int | None = None, ): """Initializes the instance. @@ -152,6 +154,17 @@ def __init__( trainer's own callbacks. Use `keras_utils.MetricsCallback` to report metrics, because Keras does not return changes made to `logs` in `on_test_end`. + verbose: The verbosity passed to `model.fit` and `model.evaluate`. "auto" + (the Keras default) and 1 show a progress bar that updates after every + step, 2 prints one line per epoch in `fit` and a single line in + `evaluate`, and 0 prints nothing. + steps_per_execution: If set, the number of training steps that each `jit` + call runs, which reduces the per-step host work of dispatching steps and + running batch-level callbacks. Batch-level callbacks run once per call + and see the logs of its last step. `steps_per_loop` must be divisible by + it. It applies to `train` and `train_and_evaluate`, including the + validation pass that `train_and_evaluate` runs inside `fit`. If None, + the model keeps the value that its `compile` call set. resume_training_launcher: Launcher to resume training. enable_xmanager_measurements: Whether to enable XManager measurements. enable_autoxprof: Whether to enable AutoXprof. @@ -164,6 +177,17 @@ def __init__( if model_dir is None: model_dir = "/tmp" + if steps_per_execution is not None: + if steps_per_execution < 1: + raise ValueError( + f"`steps_per_execution` ({steps_per_execution}) must be at least 1." + ) + if steps_per_loop % steps_per_execution != 0: + raise ValueError( + f"`steps_per_loop` ({steps_per_loop}) must be divisible by " + f"`steps_per_execution` ({steps_per_execution})." + ) + # This should be set before any layers are constructed and this is a # fallback in case the trainer binary doesn't already do this. if ( @@ -180,6 +204,7 @@ def __init__( self._steps_per_eval = steps_per_eval self._continuous_eval_timeout = continuous_eval_timeout self._steps_per_loop = steps_per_loop + self._steps_per_execution = steps_per_execution self._marker_path = os.path.join( model_dir, core.TRAINING_COMPLETE_MARKER_FILE ) @@ -188,6 +213,7 @@ def __init__( self._checkpoint_save_interval_epochs = checkpoint_save_interval_epochs self._custom_train_callbacks = list(train_callbacks) self._custom_eval_callbacks = list(eval_callbacks) + self._verbose = verbose if checkpoint_version is not None: self._checkpoint_version = keras_utils.CheckpointVersion( checkpoint_version @@ -322,6 +348,17 @@ def _maybe_get_model_kws( return kws + def _maybe_set_steps_per_execution(self, model: keras.Model): + """Sets steps_per_execution and resets cached execution functions.""" + if self._steps_per_execution is None: + return + model.steps_per_execution = self._steps_per_execution + # Clear functions so they are rebuilt by fit with the new value, mirroring + # Keras Trainer.compile (third_party/py/keras/src/trainers/trainer.py:223). + model.train_function = None + model.test_function = None + model.predict_function = None + def _val_logs( self, history: Mapping[str, Any], @@ -352,12 +389,15 @@ def train(self, task: KerasTask) -> core.Logs: **self._maybe_get_model_kws(task, dataset, val=False) ) + self._maybe_set_steps_per_execution(model) + history = model.fit( dataset, epochs=self._train_epochs + 1, steps_per_epoch=self._steps_per_loop, callbacks=self.train_callbacks, initial_epoch=1, + verbose=self._verbose, ) model.summary(print_fn=logging.info) @@ -387,6 +427,7 @@ def evaluate(self, task: KerasTask) -> core.Logs: steps=self._steps_per_eval, callbacks=self.eval_callbacks, return_dict=True, + verbose=self._verbose, ) epoch_dt = time.time() - epoch_start_time steps_per_second = self._steps_per_eval / epoch_dt # pyrefly: ignore[unsupported-operation] @@ -399,6 +440,7 @@ def evaluate(self, task: KerasTask) -> core.Logs: dataset, steps=self._steps_per_eval, callbacks=self.eval_callbacks, + verbose=self._verbose, ) def train_and_evaluate(self, task: KerasTask) -> core.Logs: @@ -417,6 +459,8 @@ def train_and_evaluate(self, task: KerasTask) -> core.Logs: **self._maybe_get_model_kws(task, train_dataset, val=False) ) + self._maybe_set_steps_per_execution(model) + history = model.fit( train_dataset, validation_data=eval_dataset, @@ -426,6 +470,7 @@ def train_and_evaluate(self, task: KerasTask) -> core.Logs: validation_steps=None, callbacks=self.train_callbacks, initial_epoch=1, + verbose=self._verbose, ) model.summary(print_fn=logging.info) @@ -514,6 +559,7 @@ def on_test_begin(self, logs: Mapping[str, Any] | None = None): steps=self._steps_per_eval, callbacks=callbacks, return_dict=True, + verbose=self._verbose, ) logging.info( diff --git a/recml/core/training/keras_trainer_test.py b/recml/core/training/keras_trainer_test.py index f396e85..a1f5b33 100644 --- a/recml/core/training/keras_trainer_test.py +++ b/recml/core/training/keras_trainer_test.py @@ -158,6 +158,198 @@ def test_keras_task_and_trainer( [7.0] * len(history.history["loss"]), ) + def test_verbose_passed_to_keras(self): + original_fit = keras.Model.fit + original_evaluate = keras.Model.evaluate + + with mock.patch.object( + keras.Model, "fit", autospec=True, side_effect=original_fit + ) as mock_fit, mock.patch.object( + keras.Model, "evaluate", autospec=True, side_effect=original_evaluate + ) as mock_evaluate: + trainer = keras_trainer.KerasTrainer( + model_dir=self.create_tempdir().full_path, + train_steps=2, + steps_per_loop=2, + steps_per_eval=1, + verbose=2, + continuous_eval_timeout=5, + distribution=keras.distribution.DataParallel() + if keras.backend.backend() == "jax" + else None, + ) + experiment = core.Experiment(_KerasTask(), trainer) + + mock_fit.reset_mock() + mock_evaluate.reset_mock() + core.run_experiment(experiment, core.Trainer.Mode.TRAIN) + mock_fit.assert_called() + self.assertEqual(mock_fit.call_args.kwargs.get("verbose"), 2) + + mock_fit.reset_mock() + mock_evaluate.reset_mock() + core.run_experiment(experiment, core.Trainer.Mode.EVAL) + mock_evaluate.assert_called() + self.assertEqual(mock_evaluate.call_args.kwargs.get("verbose"), 2) + + mock_fit.reset_mock() + mock_evaluate.reset_mock() + core.run_experiment(experiment, core.Trainer.Mode.TRAIN_AND_EVAL) + mock_fit.assert_called() + self.assertEqual(mock_fit.call_args.kwargs.get("verbose"), 2) + + if keras.backend.backend() == "jax": + keras.backend.clear_session() + mock_fit.reset_mock() + mock_evaluate.reset_mock() + core.run_experiment(experiment, core.Trainer.Mode.CONTINUOUS_EVAL) + mock_evaluate.assert_called() + self.assertEqual(mock_evaluate.call_args.kwargs.get("verbose"), 2) + + def test_default_verbose_passed_to_keras(self): + original_fit = keras.Model.fit + original_evaluate = keras.Model.evaluate + + with mock.patch.object( + keras.Model, "fit", autospec=True, side_effect=original_fit + ) as mock_fit, mock.patch.object( + keras.Model, "evaluate", autospec=True, side_effect=original_evaluate + ) as mock_evaluate: + trainer = keras_trainer.KerasTrainer( + model_dir=self.create_tempdir().full_path, + train_steps=2, + steps_per_loop=2, + steps_per_eval=1, + continuous_eval_timeout=5, + distribution=keras.distribution.DataParallel() + if keras.backend.backend() == "jax" + else None, + ) + experiment = core.Experiment(_KerasTask(), trainer) + + mock_fit.reset_mock() + mock_evaluate.reset_mock() + core.run_experiment(experiment, core.Trainer.Mode.TRAIN) + mock_fit.assert_called() + self.assertEqual(mock_fit.call_args.kwargs.get("verbose"), "auto") + + mock_fit.reset_mock() + mock_evaluate.reset_mock() + core.run_experiment(experiment, core.Trainer.Mode.EVAL) + mock_evaluate.assert_called() + self.assertEqual(mock_evaluate.call_args.kwargs.get("verbose"), "auto") + + mock_fit.reset_mock() + mock_evaluate.reset_mock() + core.run_experiment(experiment, core.Trainer.Mode.TRAIN_AND_EVAL) + mock_fit.assert_called() + self.assertEqual(mock_fit.call_args.kwargs.get("verbose"), "auto") + + if keras.backend.backend() == "jax": + keras.backend.clear_session() + mock_fit.reset_mock() + mock_evaluate.reset_mock() + core.run_experiment(experiment, core.Trainer.Mode.CONTINUOUS_EVAL) + mock_evaluate.assert_called() + self.assertEqual(mock_evaluate.call_args.kwargs.get("verbose"), "auto") + + @parameterized.named_parameters( + ("set", 2, 2), + ("default", None, 1), + ) + def test_steps_per_execution_passed_to_keras( + self, configured_value, expected_value + ): + model_steps_per_execution = [] + original_fit = keras.Model.fit + + def original_fit_with_record(model, *args, **kwargs): + model_steps_per_execution.append( + getattr(model, "steps_per_execution", None) + ) + return original_fit(model, *args, **kwargs) + + with mock.patch.object( + keras.Model, "fit", autospec=True, side_effect=original_fit_with_record + ) as mock_fit: + trainer = keras_trainer.KerasTrainer( + model_dir=self.create_tempdir().full_path, + train_steps=4, + steps_per_loop=4, + steps_per_eval=1, + steps_per_execution=configured_value, + distribution=keras.distribution.DataParallel() + if keras.backend.backend() == "jax" + else None, + ) + experiment = core.Experiment(_KerasTask(), trainer) + + model_steps_per_execution.clear() + core.run_experiment(experiment, core.Trainer.Mode.TRAIN) + mock_fit.assert_called() + self.assertEqual(model_steps_per_execution, [expected_value]) + + model_steps_per_execution.clear() + core.run_experiment(experiment, core.Trainer.Mode.TRAIN_AND_EVAL) + mock_fit.assert_called() + self.assertEqual(model_steps_per_execution, [expected_value]) + + def test_steps_per_execution_resets_functions(self): + model_train_functions = [] + original_fit = keras.Model.fit + task = _KerasTask() + original_create_model = task.create_model + + def mock_create_model(*args, **kwargs): + model = original_create_model(*args, **kwargs) + model.train_function = "stale_train" + model.test_function = "stale_test" + model.predict_function = "stale_predict" + return model + + def original_fit_with_record(model, *args, **kwargs): + model_train_functions.append(getattr(model, "train_function", None)) + return original_fit(model, *args, **kwargs) + + with mock.patch.object( + task, "create_model", autospec=True, side_effect=mock_create_model + ), mock.patch.object( + keras.Model, "fit", autospec=True, side_effect=original_fit_with_record + ) as mock_fit: + trainer = keras_trainer.KerasTrainer( + model_dir=self.create_tempdir().full_path, + train_steps=4, + steps_per_loop=4, + steps_per_eval=1, + steps_per_execution=2, + distribution=keras.distribution.DataParallel() + if keras.backend.backend() == "jax" + else None, + ) + experiment = core.Experiment(task, trainer) + core.run_experiment(experiment, core.Trainer.Mode.TRAIN) + mock_fit.assert_called() + self.assertEqual(model_train_functions, [None]) + + def test_steps_per_execution_validation(self): + with self.assertRaisesRegex( + ValueError, r"`steps_per_execution` \(0\) must be at least 1\." + ): + keras_trainer.KerasTrainer( + steps_per_loop=4, + steps_per_execution=0, + ) + + with self.assertRaisesRegex( + ValueError, + r"`steps_per_loop` \(5\) must be divisible by `steps_per_execution`" + r" \(2\)\.", + ): + keras_trainer.KerasTrainer( + steps_per_loop=5, + steps_per_execution=2, + ) + def _run_eval_with_callbacks(self, model_dir: str, mode: core.Trainer.Mode): if keras.backend.backend() != "jax": self.skipTest("`EpochSummaryCallback` is only used on the Jax backend.") diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py index e7315cf..7d90a28 100644 --- a/recml/core/utils/keras_utils.py +++ b/recml/core/utils/keras_utils.py @@ -1744,6 +1744,20 @@ def __init__( # asynchronously. self.async_safe = True + def on_train_batch_begin( + self, batch: int, logs: dict[str, Any] | None = None + ): + # Skips the per-batch bookkeeping of the inherited hook; + # see `on_train_batch_end`. + del batch, logs + + def on_train_batch_end(self, batch: int, logs: dict[str, Any] | None = None): + # The inherited `TensorBoard.on_train_batch_end` calls `summary.scalar` for + # every log on every step, even with `update_freq="epoch"`, when no summary + # writer records them. This callback only reports epoch metrics, and epoch + # summaries and `steps_per_second` don't use the batch hooks. + del batch, logs + def _get_or_create_writer( self, key: str, log_dir: str ) -> tf.summary.SummaryWriter: diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py index 182e096..0ef3f2e 100644 --- a/recml/core/utils/keras_utils_test.py +++ b/recml/core/utils/keras_utils_test.py @@ -923,6 +923,23 @@ def test_epoch_orbax_checkpoint_and_restore_callback_saves_at_epoch_0( checkpoint_manager.wait_until_finished() self.assertTrue(os.path.exists(os.path.join(checkpoint_dir, "0"))) + def test_epoch_summary_callback_batch_hooks_are_noops(self): + log_dir = self.create_tempdir().full_path + with mock.patch.object(keras_utils.jax, "process_index", return_value=0): + callback = keras_utils.EpochSummaryCallback( + log_dir=log_dir, + steps_per_epoch=10, + write_steps_per_second=False, + ) + + with mock.patch.object( + callback.summary, "scalar", autospec=True + ) as mock_scalar: + logs = {"loss": 0.5} + callback.on_train_batch_begin(0, logs=logs) + callback.on_train_batch_end(0, logs=logs) + mock_scalar.assert_not_called() + def test_epoch_summary_callback_writes_on_process_0(self): log_dir = self.create_tempdir().full_path dummy_inputs = _create_dummy_inputs()