Migrate hparams plugin to Keras 3 and remove Keras 2 from requirements.txt (#6759)

We must still install keras 2 as part of CI in order to successfully run
graph plugin tests.

Author: bmd3k

.github/workflows/ci.yml

@@ -78,13 +78,8 @@ jobs:
         run: |
           python -m pip install -U pip
           pip install "${TENSORFLOW_VERSION}"
+          pip install "${TF_KERAS_VERSION}"
         if: matrix.tf_version_id != 'notf'
-      # Replace the `tf-keras` with `tf-keras-nightly` in the requirements.txt
-      # to avoid incompatibilities when using alongside `tf-nightly`.
-      # TODO: Remove this after migrating to Keras 3.
-      - name: 'Make user to use tf-keras-nightly with tf-nightly'
-        run: |
-          sed -i "s/^tf-keras.*/${TF_KERAS_VERSION}/g" ./tensorboard/pip_package/requirements.txt
       - name: 'Install Python dependencies'
         run: |
           python -m pip install -U pip

tensorboard/pip_package/requirements.txt

@@ -33,7 +33,4 @@ setuptools >= 41.0.0  # Note: provides pkg_resources as well as setuptools
 # requirement, and likely this will not disrupt existing users of the package.
 six > 1.9
 tensorboard-data-server >= 0.7.0, < 0.8.0
-# Stay on Keras 2 for now: https://github.com/keras-team/keras/issues/18467.
-# TODO: Remove this after migrating to Keras 3.
-tf-keras >= 2.15.0
 werkzeug >= 1.0.1

tensorboard/plugins/hparams/_keras.py

@@ -24,15 +24,8 @@
 from tensorboard.plugins.hparams import summary
 from tensorboard.plugins.hparams import summary_v2
 
-# Stay on Keras 2 for now: https://github.com/keras-team/keras/issues/18467.
-version_fn = getattr(tf.keras, "version", None)
-if version_fn and version_fn().startswith("3."):
-    import tf_keras as keras  # Keras 2
-else:
-    keras = tf.keras  # Keras 2
 
-
-class Callback(keras.callbacks.Callback):
+class Callback(tf.keras.callbacks.Callback):
     """Callback for logging hyperparameters to TensorBoard.
 
     NOTE: This callback only works in TensorFlow eager mode.
@@ -41,7 +34,7 @@ class Callback(keras.callbacks.Callback):
     def __init__(self, writer, hparams, trial_id=None):
         """Create a callback for logging hyperparameters to TensorBoard.
 
-        As with the standard `keras.callbacks.TensorBoard` class, each
+        As with the standard `tf.keras.callbacks.TensorBoard` class, each
         callback object is valid for only one call to `model.fit`.
 
         Args:

tensorboard/plugins/hparams/_keras_test.py

@@ -18,22 +18,14 @@
 from unittest import mock
 
 from google.protobuf import text_format
+import numpy as np
 import tensorflow as tf
 
 from tensorboard.plugins.hparams import _keras
 from tensorboard.plugins.hparams import metadata
 from tensorboard.plugins.hparams import plugin_data_pb2
 from tensorboard.plugins.hparams import summary_v2 as hp
 
-# Stay on Keras 2 for now: https://github.com/keras-team/keras/issues/18467.
-version_fn = getattr(tf.keras, "version", None)
-if version_fn and version_fn().startswith("3."):
-    import tf_keras as keras  # Keras 2
-else:
-    keras = tf.keras  # Keras 2
-
-tf.compat.v1.enable_eager_execution()
-
 
 class CallbackTest(tf.test.TestCase):
     def setUp(self):
@@ -46,12 +38,12 @@ def _initialize_model(self, writer):
             "optimizer": "adam",
             HP_DENSE_NEURONS: 8,
         }
-        self.model = keras.models.Sequential(
+        self.model = tf.keras.models.Sequential(
             [
-                keras.layers.Dense(
+                tf.keras.layers.Dense(
                     self.hparams[HP_DENSE_NEURONS], input_shape=(1,)
                 ),
-                keras.layers.Dense(1, activation="sigmoid"),
+                tf.keras.layers.Dense(1, activation="sigmoid"),
             ]
         )
         self.model.compile(loss="mse", optimizer=self.hparams["optimizer"])
@@ -69,7 +61,11 @@ def mock_time():
         initial_time = mock_time.time
         with mock.patch("time.time", mock_time):
             self._initialize_model(writer=self.logdir)
-            self.model.fit(x=[(1,)], y=[(2,)], callbacks=[self.callback])
+            self.model.fit(
+                x=tf.constant([(1,)]),
+                y=tf.constant([(2,)]),
+                callbacks=[self.callback],
+            )
         final_time = mock_time.time
 
         files = os.listdir(self.logdir)
@@ -142,7 +138,11 @@ def test_explicit_writer(self):
             filename_suffix=".magic",
         )
         self._initialize_model(writer=writer)
-        self.model.fit(x=[(1,)], y=[(2,)], callbacks=[self.callback])
+        self.model.fit(
+            x=tf.constant([(1,)]),
+            y=tf.constant([(2,)]),
+            callbacks=[self.callback],
+        )
 
         files = os.listdir(self.logdir)
         self.assertEqual(len(files), 1, files)
@@ -158,15 +158,27 @@ def test_non_eager_failure(self):
             with self.assertRaisesRegex(
                 RuntimeError, "only supported in TensorFlow eager mode"
             ):
-                self.model.fit(x=[(1,)], y=[(2,)], callbacks=[self.callback])
+                self.model.fit(
+                    x=np.ones((10, 10)),
+                    y=np.ones((10, 10)),
+                    callbacks=[self.callback],
+                )
 
     def test_reuse_failure(self):
         self._initialize_model(writer=self.logdir)
-        self.model.fit(x=[(1,)], y=[(2,)], callbacks=[self.callback])
+        self.model.fit(
+            x=tf.constant([(1,)]),
+            y=tf.constant([(2,)]),
+            callbacks=[self.callback],
+        )
         with self.assertRaisesRegex(
             RuntimeError, "cannot be reused across training sessions"
         ):
-            self.model.fit(x=[(1,)], y=[(2,)], callbacks=[self.callback])
+            self.model.fit(
+                x=tf.constant([(1,)]),
+                y=tf.constant([(2,)]),
+                callbacks=[self.callback],
+            )
 
     def test_invalid_writer(self):
         with self.assertRaisesRegex(