Adapt to the multi-backend Keras (#928)

Author: haifeng-jin

.github/workflows/actions.yml

@@ -12,7 +12,7 @@ permissions:
 
 jobs:
   build:
-    name: Run tests
+    name: Test the code with tf.keras
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
@@ -33,20 +33,25 @@ jobs:
       - name: Install dependencies
         run: |
           pip install -e ".[tensorflow-cpu,tests]" --progress-bar off --upgrade
+          pip install jax[cpu]
       - name: Test with pytest
         run: |
           pytest --cov=keras_tuner --cov-report xml:coverage.xml
       - name: Codecov
         uses: codecov/codecov-action@v3
-  compatibility:
-    name: Check compatibility with TF 2.0
+  multibackend:
+    name: Test the code with Keras Core
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [tensorflow, jax, torch]
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
-      - name: Set up Python 3.7
+      - name: Set up Python 3.10
         uses: actions/setup-python@v4
         with:
-          python-version: '3.7'
+          python-version: '3.10'
       - name: Get pip cache dir
         id: pip-cache
         run: |
@@ -59,12 +64,15 @@ jobs:
           key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
       - name: Install dependencies
         run: |
-          pip install protobuf==3.20.1
-          pip install "tensorflow<2.1"
+          pip install -e ".[tensorflow-cpu,tests]" --progress-bar off --upgrade
+          pip install torch>=2.0.1+cpu --progress-bar off
+          pip install jax[cpu] --progress-bar off
           pip install -e ".[tensorflow-cpu,tests]" --progress-bar off --upgrade
       - name: Test with pytest
+        env:
+          KERAS_BACKEND: ${{ matrix.backend }}
         run: |
-          pytest keras_tuner/integration_tests/legacy_import_test.py
+          pytest
   format:
     name: Check the code format
     runs-on: ubuntu-latest
@@ -115,7 +123,7 @@ jobs:
       - name: Run the guides
         run: bash shell/run_guides.sh
   deploy:
-    needs: [build, format,compatibility]
+    needs: [build, format,multibackend]
     if: github.event_name == 'release' && github.event.action == 'created'
     runs-on: ubuntu-latest
     steps:
@@ -127,10 +135,10 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip setuptools wheel twine
-    - name: Build and publish
-      env:
-        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
-        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
-      run: |
-        python setup.py sdist bdist_wheel
-        twine upload dist/*
+    - name: Build a binary wheel and a source tarball
+      run: >-
+        python pip_build.py
+    - name: Publish distribution to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        password: ${{ secrets.PYPI_PASSWORD  }}

.github/workflows/nightly.yml

@@ -1,3 +1,4 @@
+import argparse
 import glob
 import os
 import shutil
@@ -6,38 +7,80 @@
 
 package = "keras_tuner"
 build_directory = "build"
-if os.path.exists(build_directory):
-    raise ValueError(f"Directory already exists: {build_directory}")
-
-# Copy sources (`keras_tuner/` directory and setup files) to build directory
-os.mkdir(build_directory)
-shutil.copytree(package, os.path.join(build_directory, package))
-shutil.copy("setup.py", os.path.join(f"{build_directory}", "setup.py"))
-shutil.copy("setup.cfg", os.path.join(f"{build_directory}", "setup.cfg"))
-os.chdir(build_directory)
-
-# Restructure the codebase so that source files live in `keras_tuner/src`
-namex.convert_codebase(package, code_directory="src")
-# Generate API __init__.py files in `keras_tuner/`
-namex.generate_api_files(package, code_directory="src", verbose=True)
-
-# Make sure to export the __version__ string
-from keras_tuner.src import __version__  # noqa: E402
-
-with open(os.path.join(package, "__init__.py")) as f:
-    init_contents = f.read()
-with open(os.path.join(package, "__init__.py"), "w") as f:
-    f.write(init_contents + "\n\n" + f'__version__ = "{__version__}"\n')
-
-# Build the package
-os.system("python3 -m build")
-
-# Save the dist files generated by the build process
-os.chdir("..")
-if not os.path.exists("dist"):
-    os.mkdir("dist")
-for filename in glob.glob(os.path.join(build_directory, "dist", "*.*")):
-    shutil.copy(filename, "dist")
-
-# Clean up: remove the build directory (no longer needed)
-shutil.rmtree(build_directory)
+dist_directory = "dist"
+to_copy = [
+    "setup.py",
+    "setup.cfg",
+    "README.md",
+]
+
+
+def build():
+    if os.path.exists(build_directory):
+        raise ValueError(f"Directory already exists: {build_directory}")
+
+    whl_path = None
+    try:
+        # Copy sources (`keras_tuner/` directory and setup files) to build
+        # directory
+        os.mkdir(build_directory)
+        shutil.copytree(package, os.path.join(build_directory, package))
+        for fname in to_copy:
+            shutil.copy(fname, os.path.join(f"{build_directory}", fname))
+        os.chdir(build_directory)
+
+        # Restructure the codebase so that source files live in
+        # `keras_tuner/src`
+        namex.convert_codebase(package, code_directory="src")
+
+        # Generate API __init__.py files in `keras_tuner/`
+        namex.generate_api_files(package, code_directory="src", verbose=True)
+
+        # Make sure to export the __version__ string
+        from keras_tuner.src import __version__  # noqa: E402
+
+        with open(os.path.join(package, "__init__.py")) as f:
+            init_contents = f.read()
+        with open(os.path.join(package, "__init__.py"), "w") as f:
+            f.write(init_contents + "\n\n" + f'__version__ = "{__version__}"\n')
+
+        # Build the package
+        os.system("python3 -m build")
+
+        # Save the dist files generated by the build process
+        os.chdir("..")
+        if not os.path.exists(dist_directory):
+            os.mkdir(dist_directory)
+        for filename in glob.glob(
+            os.path.join(build_directory, dist_directory, "*.*")
+        ):
+            shutil.copy(filename, dist_directory)
+
+        # Find the .whl file path
+        for fname in os.listdir(dist_directory):
+            if __version__ in fname and fname.endswith(".whl"):
+                whl_path = os.path.abspath(os.path.join(dist_directory, fname))
+        print(f"Build successful. Wheel file available at {whl_path}")
+
+    finally:
+        # Clean up: remove the build directory (no longer needed)
+        shutil.rmtree(build_directory)
+    return whl_path
+
+
+def install_whl(whl_fpath):
+    print("Installing wheel file.")
+    os.system(f"pip3 install {whl_fpath} --force-reinstall --no-dependencies")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--install",
+        action="store_true",
+        help="Whether to install the generated wheel file.",
+    )
+    args = parser.parse_args()
+    whl_path = build()
+    if whl_path and args.install:
+        install_whl(whl_path)

build.py

@@ -29,8 +29,4 @@
 from keras_tuner.tuners import Hyperband
 from keras_tuner.tuners import RandomSearch
 from keras_tuner.tuners import SklearnTuner
-from keras_tuner.utils import check_tf_version
-
-check_tf_version()
-
-__version__ = "1.3.5"
+from keras_tuner.version import __version__

keras_tuner/__init__.py

@@ -13,31 +13,57 @@
 # limitations under the License.
 
 
-import tensorflow as tf
-from tensorflow import keras
-from tensorflow.keras import layers
-
 from keras_tuner.api_export import keras_tuner_export
-
-try:
-    from tensorflow.keras.layers.experimental import (  # isort:skip
-        preprocessing,
-    )  # pytype: disable=import-error
-except ImportError:  # pragma: no cover
-    preprocessing = None  # pragma: no cover
-
+from keras_tuner.backend import keras
+from keras_tuner.backend import ops
+from keras_tuner.backend import random
+from keras_tuner.backend.keras import layers
 from keras_tuner.engine import hypermodel
 
 # dict of functions that create layers for transforms.
 # Each function takes a factor (0 to 1) for the strength
 # of the transform.
-if preprocessing is not None:
-    TRANSFORMS = {
-        "translate_x": lambda x: preprocessing.RandomTranslation(x, 0),
-        "translate_y": lambda y: preprocessing.RandomTranslation(0, y),
-        "rotate": preprocessing.RandomRotation,
-        "contrast": preprocessing.RandomContrast,
-    }
+TRANSFORMS = {
+    "translate_x": lambda x: layers.RandomTranslation(x, 0),
+    "translate_y": lambda y: layers.RandomTranslation(0, y),
+    "rotate": layers.RandomRotation,
+    "contrast": layers.RandomContrast,
+}
+
+
+@keras.saving.register_keras_serializable(package="keras_tuner")
+class RandAugment(keras.layers.Layer):
+    """A single RandAugment layer."""
+
+    def __init__(self, layers, factors):
+        super().__init__()
+        self.layers = layers
+        self.factors = factors
+
+    def call(self, inputs):
+        x = inputs
+        batch_size = ops.shape(x)[0]
+        # selection tensor determines operation for each sample.
+        selection = ops.cast(
+            random.uniform((batch_size, 1, 1, 1), maxval=len(self.layers)),
+            dtype="int32",
+        )
+
+        for i, layer in enumerate(self.layers):
+            factor = self.factors[i]
+            if factor == 0:
+                continue
+            transform_layer = TRANSFORMS[layer](factor)
+            x_trans = transform_layer(x)
+
+            # For each sample, apply the transform if and only if
+            # selection matches the transform index `i`
+            x = ops.where(ops.equal(i, selection), x_trans, x)
+
+        return x
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
 
 
 @keras_tuner_export("keras_tuner.applications.HyperImageAugment")
@@ -137,12 +163,6 @@ def __init__(
         augment_layers=3,
         **kwargs,
     ):
-        if preprocessing is None:
-            raise ImportError(
-                "HyperImageAugment requires tensorflow>=2.3.0, "
-                f"but the current version is {tf.__version__}."
-            )
-
         if input_shape is None and input_tensor is None:
             raise ValueError(
                 "You must specify either `input_shape` or `input_tensor`."
@@ -208,27 +228,14 @@ def _build_randaug_layers(self, inputs, hp):
         )
         x = inputs
         for _ in range(augment_layers):
-            # selection tensor determines operation for each sample.
-            batch_size = tf.shape(x)[0]
-            selection = tf.random.uniform(
-                [batch_size, 1, 1, 1],
-                maxval=len(self.transforms),
-                dtype="int32",
-            )
-
+            factors = []
             for i, (transform, (f_min, f_max)) in enumerate(self.transforms):
                 # Factor for each transform is determined per each trial.
                 factor = hp.Float(
                     f"factor_{transform}", f_min, f_max, default=f_min
                 )
-                if factor == 0:
-                    continue
-                transform_layer = TRANSFORMS[transform](factor)
-                x_trans = transform_layer(x)
-
-                # For each sample, apply the transform if and only if
-                # selection matches the transform index `i`
-                x = tf.where(tf.equal(i, selection), x_trans, x)
+                factors.append(factor)
+            x = RandAugment([layer for layer, _ in self.transforms], factors)(x)
         return x
 
     def _build_fixedaug_layers(self, inputs, hp):