From 8a99668b3806b6f6a8b2493ce86f6a5cdae96ad4 Mon Sep 17 00:00:00 2001 From: Brian Dubois Date: Thu, 15 Feb 2024 17:54:07 -0500 Subject: [PATCH] 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. --- .github/workflows/ci.yml | 7 +--- tensorboard/pip_package/requirements.txt | 3 -- tensorboard/plugins/hparams/_keras.py | 11 +----- tensorboard/plugins/hparams/_keras_test.py | 46 ++++++++++++++-------- 4 files changed, 32 insertions(+), 35 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9f9aa0b7f3..2a1404c98a 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/tensorboard/pip_package/requirements.txt b/tensorboard/pip_package/requirements.txt index 5ebf221504..5c259439f2 100644 --- a/tensorboard/pip_package/requirements.txt +++ b/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 diff --git a/tensorboard/plugins/hparams/_keras.py b/tensorboard/plugins/hparams/_keras.py index 004704dca2..dc57c09301 100644 --- a/tensorboard/plugins/hparams/_keras.py +++ b/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: diff --git a/tensorboard/plugins/hparams/_keras_test.py b/tensorboard/plugins/hparams/_keras_test.py index 42cbe30c96..d9a75de6a4 100644 --- a/tensorboard/plugins/hparams/_keras_test.py +++ b/tensorboard/plugins/hparams/_keras_test.py @@ -18,6 +18,7 @@ from unittest import mock from google.protobuf import text_format +import numpy as np import tensorflow as tf from tensorboard.plugins.hparams import _keras @@ -25,15 +26,6 @@ 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(