Support TensorFlow 2.11.0 (#287)

* Support TensorFlow 2.11.0 and make it the minimum supported version * Clean up the serialization helper utilities
adriangb · Dec 12, 2022 · 32fc6af · 32fc6af
1 parent eb88be8
commit 32fc6af
Show file tree

Hide file tree

Showing 9 changed files with 135 additions and 150 deletions.
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -8,6 +8,10 @@ on:
     # run every day at midnight
     - cron: "0 0 * * *"
 
+defaults:
+  run:
+    shell: bash
+
 jobs:
   Linting:
     runs-on: ubuntu-latest
@@ -30,6 +34,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.7", "3.8", "3.9", "3.10"]
+      fail-fast: false
 
     steps:
       - uses: actions/checkout@v2
@@ -56,53 +61,54 @@ jobs:
 
       - uses: codecov/codecov-action@v1
 
-  # TestDev:
-  #   name: Ubuntu / Python ${{ matrix.python-version }} / TensorFlow Nightly / Scikit-Learn Nightly
-  #   runs-on: ubuntu-latest
-  #   strategy:
-  #     matrix:
-  #       python-version: ["3.10"]
-
-  #   steps:
-  #     - uses: actions/checkout@v2
-
-  #     - name: Set up Python ${{ matrix.python-version }}
-  #       uses: actions/setup-python@v2
-  #       with:
-  #         python-version: ${{ matrix.python-version }}
-
-  #     - name: Install and Set Up Poetry
-  #       run: |
-  #         python -m pip install --upgrade poetry
-  #         poetry config virtualenvs.in-project true
-  #         poetry run python -m pip install --upgrade pip
-
-  #     - name: Install Dependencies
-  #       run: |
-  #         poetry remove scikit-learn
-  #         poetry install
-
-  #     - name: Install Nightly Versions
-  #       if: always()
-  #       run: |
-  #         poetry run python -m pip install -U tf-nightly
-  #         poetry run python -m pip install -U scipy
-  #         poetry run python -m pip install -U --pre --extra-index https://pypi.anaconda.org/scipy-wheels-nightly/simple scikit-learn
-
-  #     - name: Test with pytest
-  #       if: always()
-  #       run: |
-  #         poetry run python -m pip freeze
-  #         poetry run python -m pytest -v --cov=scikeras --cov-report xml --color=yes
-
-  #     - uses: codecov/codecov-action@v1
+  TestDev:
+    name: Ubuntu / Python ${{ matrix.python-version }} / TensorFlow Nightly / Scikit-Learn Nightly
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10"]
+      fail-fast: false
+
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install and Set Up Poetry
+        run: |
+          python -m pip install --upgrade poetry
+          poetry config virtualenvs.in-project true
+          poetry run python -m pip install --upgrade pip
+
+      - name: Install Dependencies
+        run: |
+          poetry remove scikit-learn
+          poetry install
+
+      - name: Install Nightly Versions
+        if: always()
+        run: |
+          poetry run python -m pip install -U tf-nightly
+          poetry run python -m pip install -U scipy
+          poetry run python -m pip install -U --pre --extra-index https://pypi.anaconda.org/scipy-wheels-nightly/simple scikit-learn
+
+      - name: Test with pytest
+        if: always()
+        run: |
+          poetry run python -m pip freeze
+          poetry run python -m pytest -v --cov=scikeras --cov-report xml --color=yes
+
+      - uses: codecov/codecov-action@v1
 
   TestOldest:
     name: Ubuntu / Python ${{ matrix.python-version }} / TF ${{ matrix.tf-version }} / Scikit-Learn ${{ matrix.sklearn-version }}
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        tf-version: [2.7.0]
+        tf-version: [2.11.0]
         python-version: ["3.7", "3.9"]
         sklearn-version: [1.0.0]
 
@@ -123,7 +129,9 @@ jobs:
       - name: Install Dependencies
         run: |
           poetry install -E tensorflow
-          poetry add "scikit-learn==${{matrix.sklearn-version}}"
+          poetry run pip install \
+            "tensorflow==${{ matrix.tf-version }}" \
+            "scikit-learn==${{matrix.sklearn-version}}"
 
       - name: Test with pytest
         run: |
@@ -139,6 +147,7 @@ jobs:
       matrix:
         os: [MacOS, Windows]  # test all OSs (except Ubuntu, which is already running other tests)
         python-version: ["3.7", "3.10"]  # test only the two extremes of supported Python versions
+      fail-fast: false
 
     steps:
       - uses: actions/checkout@v2
@@ -149,19 +158,24 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install and Set Up Poetry
-        shell: bash
         run: |
           pip install --upgrade poetry
           poetry config virtualenvs.in-project true
           poetry run python -m pip install --upgrade pip
 
       - name: Install Dependencies
+        # TF is sorta dropping support for Windows
+        # At the very least they are no longer suppporting GPU
+        # See https://github.com/tensorflow/tensorflow/releases/tag/v2.11.0
+        # and
+        # https://www.tensorflow.org/install/pip#windows-native
         run: |
-          poetry install -E tensorflow
+          poetry install
+          poetry run pip install -U tensorflow-cpu
 
       - name: Test with pytest
         run: |
-          poetry run python -m pip freeze
+          poetry run pip freeze
           poetry run python -m pytest -v --cov=scikeras --cov-report xml --color=yes
 
       - uses: codecov/codecov-action@v1
diff --git a/pyproject.toml b/pyproject.toml
@@ -27,22 +27,24 @@ license = "MIT"
 name = "scikeras"
 readme = "README.md"
 repository = "https://github.com/adriangb/scikeras"
-version = "0.9.0"
+version = "0.10.0"
 
 [tool.poetry.dependencies]
 importlib-metadata = {version = ">=3", python = "<3.8"}
 python = ">=3.7.0,<3.11.0"
 scikit-learn = ">=1.0.0"
 packaging = ">=0.21"
-tensorflow = {version = ">=2.7.0", optional = true}
-tensorflow-cpu = {version = ">=2.7.0", optional = true}
+tensorflow = {version = ">=2.11.0", optional = true}
+tensorflow-cpu = {version = ">=2.11.0", optional = true}
+# https://github.com/grpc/grpc/issues/31492
+grpcio = { version = "<1.50.0", markers = "python_version >= '3.10' and sys_platform == 'darwin'" }
 
 [tool.poetry.extras]
 tensorflow = ["tensorflow"]
 tensorflow-cpu = ["tensorflow-cpu"]
 
 [tool.poetry.dev-dependencies]
-tensorflow = ">=2.7.0"
+tensorflow = ">=2.11.0"
 coverage = {extras = ["toml"], version = ">=6.4.2"}
 insipid-sphinx-theme = ">=0.3.2"
 ipykernel = ">=6.15.1"

diff --git a/scikeras/_saving_utils.py b/scikeras/_saving_utils.py
@@ -1,81 +1,57 @@
 import os
+import shutil
 import tarfile
 import tempfile
 
+from contextlib import contextmanager
 from io import BytesIO
-from types import MethodType
-from typing import Any, Callable, Dict, Hashable, Iterable, List, Tuple
+from typing import Any, Callable, ContextManager, Dict, Hashable, Iterator, List, Tuple
+from uuid import uuid4
 
 import numpy as np
 import tensorflow.keras as keras
 
-from tensorflow import Variable
 from tensorflow import io as tf_io
 from tensorflow.keras.models import load_model
 
 
-def _get_temp_folder() -> str:
+@contextmanager
+def _get_temp_folder() -> Iterator[str]:
     if os.name == "nt":
         # the RAM-based filesystem is not fully supported on
         # Windows yet, we save to a temp folder on disk instead
-        return tempfile.mkdtemp()
+        tmp_dir = tempfile.mkdtemp()
+        try:
+            yield tmp_dir
+        finally:
+            shutil.rmtree(tmp_dir, ignore_errors=True)
     else:
-        return f"ram://{tempfile.mkdtemp()}"
-
-
-def _temp_create_all_weights(
-    self: keras.optimizers.Optimizer, var_list: Iterable[Variable]
-):
-    """A hack to restore weights in optimizers that use slots.
-
-    See https://tensorflow.org/api_docs/python/tf/keras/optimizers/Optimizer#slots_2
-    """
-    self._create_all_weights_orig(var_list)
-    try:
-        self.set_weights(self._restored_weights)
-    except ValueError:
-        # Weights don't match, eg. when optimizer was pickled before any training
-        # or a completely new dataset is being used right after pickling
-        pass
-    delattr(self, "_restored_weights")
-    self._create_all_weights = self._create_all_weights_orig
-
-
-def _restore_optimizer_weights(
-    optimizer: keras.optimizers.Optimizer, weights: List[np.ndarray]
-) -> None:
-    optimizer._restored_weights = weights
-    optimizer._create_all_weights_orig = optimizer._create_all_weights
-    # MethodType is used to "bind" the _temp_create_all_weights method
-    # to the "live" optimizer object
-    optimizer._create_all_weights = MethodType(_temp_create_all_weights, optimizer)
+        temp_dir = f"ram://{uuid4().hex}"
+        try:
+            yield temp_dir
+        finally:
+            for root, _, filenames in tf_io.gfile.walk(temp_dir):
+                for filename in filenames:
+                    dest = os.path.join(root, filename)
+                    tf_io.gfile.remove(dest)
 
 
 def unpack_keras_model(
-    packed_keras_model: np.ndarray, optimizer_weights: List[np.ndarray]
+    packed_keras_model: np.ndarray,
 ):
     """Reconstruct a model from the result of __reduce__"""
-    temp_dir = _get_temp_folder()
-    b = BytesIO(packed_keras_model)
-    with tarfile.open(fileobj=b, mode="r") as archive:
-        for fname in archive.getnames():
-            dest = os.path.join(temp_dir, fname)
-            tf_io.gfile.makedirs(os.path.dirname(dest))
-            with tf_io.gfile.GFile(dest, "wb") as f:
-                f.write(archive.extractfile(fname).read())
-    model: keras.Model = load_model(temp_dir)
-    for root, _, filenames in tf_io.gfile.walk(temp_dir):
-        for filename in filenames:
-            if filename.startswith("ram://"):
-                # Currently, tf.io.gfile.walk returns
-                # the entire path for the ram:// filesystem
-                dest = filename
-            else:
-                dest = os.path.join(root, filename)
-            tf_io.gfile.remove(dest)
-    if model.optimizer is not None:
-        _restore_optimizer_weights(model.optimizer, optimizer_weights)
-    return model
+    with _get_temp_folder() as temp_dir:
+        b = BytesIO(packed_keras_model)
+        with tarfile.open(fileobj=b, mode="r") as archive:
+            for fname in archive.getnames():
+                dest = os.path.join(temp_dir, fname)
+                tf_io.gfile.makedirs(os.path.dirname(dest))
+                with tf_io.gfile.GFile(dest, "wb") as f:
+                    f.write(archive.extractfile(fname).read())
+        model: keras.Model = load_model(temp_dir)
+        model.load_weights(temp_dir)
+        model.optimizer.build(model.trainable_variables)
+        return model
 
 
 def pack_keras_model(
@@ -85,47 +61,35 @@ def pack_keras_model(
     Tuple[np.ndarray, List[np.ndarray]],
 ]:
     """Support for Pythons's Pickle protocol."""
-    temp_dir = _get_temp_folder()
-    model.save(temp_dir)
-    b = BytesIO()
-    with tarfile.open(fileobj=b, mode="w") as archive:
-        for root, _, filenames in tf_io.gfile.walk(temp_dir):
-            for filename in filenames:
-                if filename.startswith("ram://"):
-                    # Currently, tf.io.gfile.walk returns
-                    # the entire path for the ram:// filesystem
-                    dest = filename
-                else:
+    with _get_temp_folder() as temp_dir:
+        model.save(temp_dir)
+        b = BytesIO()
+        with tarfile.open(fileobj=b, mode="w") as archive:
+            for root, _, filenames in tf_io.gfile.walk(temp_dir):
+                for filename in filenames:
                     dest = os.path.join(root, filename)
-                with tf_io.gfile.GFile(dest, "rb") as f:
-                    info = tarfile.TarInfo(name=os.path.relpath(dest, temp_dir))
-                    info.size = f.size()
-                    archive.addfile(tarinfo=info, fileobj=f)
-                tf_io.gfile.remove(dest)
-    b.seek(0)
-    optimizer_weights = None
-    if model.optimizer is not None:
-        optimizer_weights = model.optimizer.get_weights()
-    model_bytes = np.asarray(memoryview(b.read()))
-    return (
-        unpack_keras_model,
-        (model_bytes, optimizer_weights),
-    )
+                    with tf_io.gfile.GFile(dest, "rb") as f:
+                        info = tarfile.TarInfo(name=os.path.relpath(dest, temp_dir))
+                        info.size = f.size()
+                        archive.addfile(tarinfo=info, fileobj=f)
+                    tf_io.gfile.remove(dest)
+        b.seek(0)
+        model_bytes = np.asarray(memoryview(b.read()))
+        return (unpack_keras_model, (model_bytes,))
 
 
 def deepcopy_model(model: keras.Model, memo: Dict[Hashable, Any]) -> keras.Model:
-    _, (model_bytes, optimizer_weights) = pack_keras_model(model)
-    new_model = unpack_keras_model(model_bytes, optimizer_weights)
+    _, (model_bytes,) = pack_keras_model(model)
+    new_model = unpack_keras_model(model_bytes)
     memo[model] = new_model
     return new_model
 
 
 def unpack_keras_optimizer(
-    opt_serialized: Dict[str, Any], weights: List[np.ndarray]
+    opt_serialized: Dict[str, Any]
 ) -> keras.optimizers.Optimizer:
     """Reconstruct optimizer."""
     optimizer: keras.optimizers.Optimizer = keras.optimizers.deserialize(opt_serialized)
-    _restore_optimizer_weights(optimizer, weights)
     return optimizer
 
 
@@ -139,8 +103,7 @@ def pack_keras_optimizer(
 ]:
     """Support for Pythons's Pickle protocol in Keras Optimizers."""
     opt_serialized = keras.optimizers.serialize(optimizer)
-    weights = optimizer.get_weights()
-    return unpack_keras_optimizer, (opt_serialized, weights)
+    return unpack_keras_optimizer, (opt_serialized,)
 
 
 def unpack_keras_metric(metric_serialized: Dict[str, Any]) -> keras.metrics.Metric:

diff --git a/scikeras/utils/transformers.py b/scikeras/utils/transformers.py
@@ -47,8 +47,7 @@ def fit(self, y: np.ndarray) -> "TargetReshaper":
         self.ndim_ = y.ndim
         return self
 
-    @staticmethod
-    def transform(y: np.ndarray) -> np.ndarray:
+    def transform(self, y: np.ndarray) -> np.ndarray:
         """Makes 1D y 2D.
 
         Parameters