Integrate HParams plugin with DistributingCloudTuner

PiperOrigin-RevId: 339603283

Author: chongyouquan

src/python/tensorflow_cloud/tuner/tests/unit/tuner_test.py

@@ -25,6 +25,7 @@
 import mock
 
 import tensorflow as tf
+from tensorboard.plugins.hparams import api as hparams_api
 from tensorflow_cloud.core import deploy
 from tensorflow_cloud.core import machine_config
 from tensorflow_cloud.core import validate
@@ -125,9 +126,10 @@ def _remote_tuner(
                 objective,
                 hyperparameters,
                 study_config,
-                directory="gs://remote_dir",
+                directory=None,
                 max_trials=None
                 ):
+        directory = directory or self._remote_dir
         return tuner.DistributingCloudTuner(
             hypermodel=build_model,
             objective=objective,
@@ -457,19 +459,66 @@ def test_get_best_trials_multi_tuners(self):
         self.assertEqual(best_trials_1[0].best_step, 3)
 
     @mock.patch.object(super_tuner.Tuner, "__init__", auto_spec=True)
-    def test_add_tensorboard_callback(self, mock_super_tuner):
+    @mock.patch.object(tf.summary, "create_file_writer", auto_spec=True)
+    @mock.patch.object(hparams_api, "hparams", auto_spec=True)
+    def test_add_logging_user_specified(
+        self, mock_hparams, mock_create_file_writer, mock_super_tuner):
         remote_tuner = self._remote_tuner(None, None, self._study_config)
 
-        callbacks = [
-            tf.keras.callbacks.TensorBoard(log_dir="user_defined_path_1"),
-            tf.keras.callbacks.TensorBoard(log_dir="user_defined_path_2")]
+        callbacks = [tf.keras.callbacks.TensorBoard(
+            log_dir=remote_tuner.directory,
+            write_images=True)]
+
+        remote_tuner._add_logging(callbacks, self._test_trial)
+
+        expected_logdir = os.path.join(
+            remote_tuner.directory, self._test_trial.trial_id, "logs")
+        expected_hparams = {hparams_api.HParam(
+            "learning_rate", hparams_api.Discrete([1e-4, 1e-3, 1e-2])): 1e-4}
 
-        trial_id = "test_trial_id"
-        remote_tuner._add_tensorboard_callback(callbacks, trial_id)
         self.assertLen(callbacks, 1)
+        self.assertEqual(callbacks[0].log_dir, expected_logdir)
+        self.assertEqual(callbacks[0].write_images, True)
+        mock_create_file_writer.assert_called_once_with(expected_logdir)
+        self.assertEqual(mock_hparams.call_count, 1)
         self.assertEqual(
-            callbacks[0].log_dir,
-            os.path.join(remote_tuner.directory, trial_id, "logs"))
+            repr(mock_hparams.call_args[0][0]), repr(expected_hparams))
+
+    @mock.patch.object(super_tuner.Tuner, "__init__", auto_spec=True)
+    @mock.patch.object(tf.summary, "create_file_writer", auto_spec=True)
+    @mock.patch.object(hparams_api, "hparams", auto_spec=True)
+    def test_add_logging_not_specified(
+        self, mock_hparams, mock_create_file_writer, mock_super_tuner):
+        remote_tuner = self._remote_tuner(None, None, self._study_config)
+
+        callbacks = []
+        remote_tuner._add_logging(callbacks, self._test_trial)
+
+        expected_logdir = os.path.join(
+            remote_tuner.directory, self._test_trial.trial_id, "logs")
+
+        self.assertLen(callbacks, 1)
+        self.assertEqual(callbacks[0].log_dir, expected_logdir)
+        mock_create_file_writer.assert_not_called()
+        mock_hparams.assert_not_called()
+
+    @mock.patch.object(super_tuner.Tuner, "__init__", auto_spec=True)
+    @mock.patch.object(tf.summary, "create_file_writer", auto_spec=True)
+    @mock.patch.object(hparams_api, "hparams", auto_spec=True)
+    def test_add_logging_mismatched_dir(
+        self, mock_hparams, mock_create_file_writer, mock_super_tuner):
+        remote_tuner = self._remote_tuner(None, None, self._study_config)
+
+        callbacks = [tf.keras.callbacks.TensorBoard(
+            log_dir=os.path.join(remote_tuner.directory, "logs"))]
+
+        with self.assertRaisesRegex(
+            ValueError, "log_dir in TensorBoard callback should be "
+                        "gs://remote_dir, but was gs://remote_dir/logs"):
+            remote_tuner._add_logging(callbacks, self._test_trial)
+
+        mock_create_file_writer.assert_not_called()
+        mock_hparams.assert_not_called()
 
     @mock.patch.object(super_tuner.Tuner, "__init__", auto_spec=True)
     def test_add_model_checkpoint_callback(self, mock_super_tuner):

src/python/tensorflow_cloud/tuner/tests/unit/utils_test.py

@@ -20,6 +20,7 @@
 from kerastuner.engine import oracle as oracle_module
 from kerastuner.engine import trial as trial_module
 import tensorflow as tf
+from tensorboard.plugins.hparams import api as hparams_api
 from tensorflow_cloud.tuner.tuner import utils
 
 STUDY_CONFIG_DISCRETE = {
@@ -315,6 +316,99 @@ def test_convert_optimizer_trial_to_keras_trial(self):
         self.assertEqual(
             trial.hyperparameters.values, {"learning_rate": 0.0001})
 
+    def test_convert_hyperparams_to_hparams_choice(self):
+        hps = hp_module.HyperParameters()
+        hps.Choice("learning_rate", [1e-4, 1e-3, 1e-2])
+        hparams = utils.convert_hyperparams_to_hparams(hps)
+        expected_hparams = {
+            hparams_api.HParam("learning_rate",
+                               hparams_api.Discrete([1e-4, 1e-3, 1e-2])): 1e-4,
+        }
+        self.assertEqual(repr(hparams), repr(expected_hparams))
+
+    @parameterized.parameters(
+        ("units", 2, 16, None, hparams_api.IntInterval(2, 16), 2),
+        ("units", 32, 128, 32, hparams_api.Discrete([32, 64, 96, 128]), 32))
+    def test_convert_hyperparams_to_hparams_int(self, name, min_value,
+                                                max_value, step,
+                                                expected_domain,
+                                                expected_value):
+        hps = hp_module.HyperParameters()
+        if step:
+            hps.Int(name, min_value=min_value, max_value=max_value, step=step)
+        else:
+            hps.Int(name, min_value=min_value, max_value=max_value)
+        hparams = utils.convert_hyperparams_to_hparams(hps)
+        expected_hparams = {
+            hparams_api.HParam(name, expected_domain): expected_value,
+        }
+        self.assertEqual(repr(hparams), repr(expected_hparams))
+
+    @parameterized.parameters(
+        ("learning_rate", 0.5, 1.5, 0.25,
+         hparams_api.Discrete([0.5, 0.75, 1.0, 1.25, 1.5]), 0.5),
+        ("learning_rate", 1e-4, 1e-1, None,
+         hparams_api.RealInterval(1e-4, 1e-1), 1e-4))
+    def test_convert_hyperparams_to_hparams_float(self, name, min_value,
+                                                  max_value, step,
+                                                  expected_domain,
+                                                  expected_value):
+        hps = hp_module.HyperParameters()
+        hps.Float(name, min_value=min_value, max_value=max_value, step=step)
+        hparams = utils.convert_hyperparams_to_hparams(hps)
+        expected_hparams = {
+            hparams_api.HParam(name, expected_domain): expected_value,
+        }
+        self.assertEqual(repr(hparams), repr(expected_hparams))
+
+    def test_convert_hyperparams_to_hparams_multi_float(self):
+        hps = hp_module.HyperParameters()
+        hps.Float("theta", min_value=0.0, max_value=1.57)
+        hps.Float("r", min_value=0.0, max_value=1.0)
+        hparams = utils.convert_hyperparams_to_hparams(hps)
+        expected_hparams = {
+            hparams_api.HParam("r", hparams_api.RealInterval(0.0, 1.0)): 0.0,
+            hparams_api.HParam("theta",
+                               hparams_api.RealInterval(0.0, 1.57)): 0.0,
+        }
+        hparams_repr_list = [repr(hparams[x]) for x in hparams.keys()]
+        expected_hparams_repr_list = [
+            repr(expected_hparams[x]) for x in expected_hparams.keys()
+        ]
+        self.assertCountEqual(hparams_repr_list, expected_hparams_repr_list)
+
+    def test_convert_hyperparams_to_hparams_boolean(self):
+        hps = hp_module.HyperParameters()
+        hps.Boolean("has_beta")
+        hparams = utils.convert_hyperparams_to_hparams(hps)
+        expected_hparams = {
+            hparams_api.HParam("has_beta", hparams_api.Discrete([True, False])):
+                False,
+        }
+        self.assertEqual(repr(hparams), repr(expected_hparams))
+
+    @parameterized.parameters(
+        ("beta", 0.1),
+        ("type", "WIDE_AND_DEEP"),
+        ("num_layers", 2))
+    def test_convert_hyperparams_to_hparams_fixed(self, name, value):
+        hps = hp_module.HyperParameters()
+        hps.Fixed(name, value)
+        hparams = utils.convert_hyperparams_to_hparams(hps)
+        expected_hparams = {
+            hparams_api.HParam(name, hparams_api.Discrete([value])): value,
+        }
+        self.assertEqual(repr(hparams), repr(expected_hparams))
+
+    def test_convert_hyperparams_to_hparams_fixed_bool(self):
+        hps = hp_module.HyperParameters()
+        hps.Fixed("condition", True)
+        hparams = utils.convert_hyperparams_to_hparams(hps)
+        expected_hparams = {
+            hparams_api.HParam("condition", hparams_api.Discrete([1])): 1,
+        }
+        self.assertEqual(repr(hparams), repr(expected_hparams))
+
     @parameterized.parameters(
         ("val_loss", "min",
          [oracle_module.Objective(name="val_loss", direction="min")]),

src/python/tensorflow_cloud/tuner/tuner.py

@@ -28,6 +28,7 @@
 from kerastuner.engine import tuner as tuner_module
 import tensorflow as tf
 
+from tensorboard.plugins.hparams import api as hparams_api
 from tensorflow_cloud.core import deploy
 from tensorflow_cloud.core import machine_config
 from tensorflow_cloud.core import validate
@@ -492,8 +493,8 @@ def __init__(
         super(DistributingCloudTuner, self,).__init__(
             oracle=oracle, hypermodel=hypermodel, **kwargs
         )
-        # If study id is not provided cloud_oracle creates ones. Setting the
-        # study_id based on cloud oracles logic to ensure they are the same.
+        # If study_id is not provided, CloudOracle creates one. Setting the
+        # study_id to what CloudOracle generates, to ensure they are the same.
         self._study_id = oracle.study_id
         self.directory = directory
 
@@ -519,16 +520,15 @@ def run_trial(self, trial, *fit_args, **fit_kwargs):
         callbacks = fit_kwargs.pop("callbacks", [])
         callbacks = self._deepcopy_callbacks(callbacks)
 
-        # Note run_trial does not use `TunerCallback` calls, since
+        # Note: run_trial does not use `TunerCallback` calls, since
         # training is performed on AI Platform training remotely.
 
-        # Creating a tensorboard callback with log-dir path specific for this
-        # trail_id. The tensorboard logs are used for passing metrics back from
-        # remote execution.
-        self._add_tensorboard_callback(callbacks, trial.trial_id)
+        # Handle TensorBoard/hyperparameter logging here. The TensorBoard
+        # logs are used for passing metrics back from remote execution.
+        self._add_logging(callbacks, trial)
 
         # Creating a save_model checkpoint callback with a saved model file path
-        # specific to this trial, this is to prevent different trials from
+        # specific to this trial. This is to prevent different trials from
         # overwriting each other.
         self._add_model_checkpoint_callback(
             callbacks, trial.trial_id)
@@ -605,7 +605,9 @@ def run_trial(self, trial, *fit_args, **fit_kwargs):
         if not google_api_client.wait_for_api_training_job_completion(
             job_id, self._project_id):
             raise RuntimeError(
-                "AIP Training job failed, see logs for details at https://console.cloud.google.com/ai-platform/jobs/{}/charts/cpu?project={}"  # pylint: disable=line-too-long
+                "AIP Training job failed, see logs for details at "
+                "https://console.cloud.google.com/ai-platform/jobs/"
+                "{}/charts/cpu?project={}"
                 .format(job_id, self._project_id))
 
         # Retrieve and report any remaining metrics
@@ -657,7 +659,7 @@ def _get_remote_training_metrics(
         self,
         log_reader,
         partial_epoch_metrics: Dict[Text, float]
-        )-> _TrainingMetrics:
+        ) -> _TrainingMetrics:
         """Retrieves delta epoch metrics from tensorboard logs since last run.
 
         This method reports any complete epoch metrics that are available since
@@ -683,9 +685,9 @@ def _get_remote_training_metrics(
         completed_epoch_metrics = []
         for event in log_reader.Load():
             for value in event.summary.value:
-                # Note tf.keras.callbacks.TensorBoard() with update_freq="epoch"
-                # logs the epoch related metrics with a "epoch_" prefix. This is
-                # not a requirement by tensorboard.
+                # Note: tf.keras.callbacks.TensorBoard.on_epoch_end() logs the
+                # epoch related metrics with a "epoch_" prefix. Please refer to
+                # https://github.com/tensorflow/tensorflow/blob/fcc4b966f1265f466e82617020af93670141b009/tensorflow/python/keras/callbacks.py#L2179 # pylint: disable=line-too-long
                 if value.tag.startswith("epoch_"):
                     metric = value.tag.replace("epoch_", "")
                     # If we have already seen this metric, this is a new epoch
@@ -708,7 +710,6 @@ def load_model(self, trial):
         raise NotImplementedError("load_model for remote run is not supported.")
 
     def save_model(self, trial_id: int, model, step: int = 0):
-
         # In remote execution models are saved automatically in Google Cloud
         # Storage (GCS) bucket hence no additional actions are needed to save
         # the model.
@@ -719,27 +720,58 @@ def _add_model_checkpoint_callback(self, callbacks, trial_id):
             filepath=self._get_model_checkpoint_dir(trial_id),
             save_freq="epoch"))
 
-    def _add_tensorboard_callback(self, callbacks, trial_id):
-        # due to https://github.com/keras-team/keras/issues/14223 multiple
-        # tensorboard callbacks are not supported. Removing user defined
-        # tf.keras.callbacks.TensorBoard callback.
+    def _add_logging(self, callbacks, trial):
+        """Add a TensorBoard callback if needed, otherwise log hyperparameters.
 
-        tf.get_logger().info(
-            "Only one tf.keras.callbacks.TensorBoard callback is allowed, removing user defined callbacks."  # pylint: disable=line-too-long
-            )
-        callbacks[:] = [
-            x for x in callbacks if x.__class__.__name__ != "TensorBoard"]
+        Note: Due to https://github.com/keras-team/keras/issues/14223, multiple
+        TensorBoard callbacks are not supported. If user specified a TensorBoard
+        callback, we treat it as an intent to log the metrics, and we shall
+        additionally log the hyperparameters as well. Otherwise, we'll add a
+        TensorBoard callback to pass back the epoch related metrics from
+        remote execution.
 
-        callbacks.append(tf.keras.callbacks.TensorBoard(
-            log_dir=self._get_tensorboard_log_dir(trial_id)))
+        Arguments:
+            callbacks: List of callbacks passed in to the search function.
+            trial: A `Trial` instance.
+        Raises:
+            ValueError: If TensorBoard callback's log_dir does not match
+            self.directory.
+        """
+
+        logdir = self._get_tensorboard_log_dir(trial.trial_id)
+        for callback in callbacks:
+            if callback.__class__.__name__ == "TensorBoard":
+                # Validate TensorBoard log_dir
+                if callback.log_dir != self.directory:
+                    # TODO(b/170687807) Switch from using .format() to f-string
+                    raise ValueError(
+                        "log_dir in TensorBoard callback should be {}, "
+                        "but was {}".format(self.directory, callback.log_dir)
+                    )
+                # Patch the log_dir
+                callback.log_dir = logdir
+                # Do hyperparameter logging here to avoid having to
+                # serialize/deserialize the hyperparameters if logged through
+                # passing hparams_api.KerasCallback to client.cloud_fit.
+                with tf.summary.create_file_writer(logdir).as_default():
+                    hparams_api.hparams(utils.convert_hyperparams_to_hparams(
+                        trial.hyperparameters))
+                # We're done here, since there should only be one TensorBoard
+                # callback
+                return
+
+        # TensorBoard callback not specified by user, add it here. The
+        # TensorBoard logs are used for passing metrics back from
+        # remote execution.
+        callbacks.append(tf.keras.callbacks.TensorBoard(log_dir=logdir))
 
-    def _get_tensorboard_log_dir(self, trial_id)-> Text:
+    def _get_tensorboard_log_dir(self, trial_id) -> Text:
         # Defining <directory>/<trial_id>/logs as log structure.
         # self._add_tensorboard_callback uses this directory structure to
         # configure the tf.keras.callbacks.TensorBoard() for each trial.
         return os.path.join(self.directory, str(trial_id), "logs")
 
-    def _get_model_checkpoint_dir(self, trial_id)->Text:
+    def _get_model_checkpoint_dir(self, trial_id) -> Text:
         # Defining <directory>/<trial_id>/checkpoint as checkpoint structure.
         # self._add_model_checkpoint_callback uses this directory structure to
         # configure the tf.keras.callbacks.ModelCheckpoint() for each trial.