DOC, TST: Wrapping of Keras models (#713)

Author: stsievert

ci/posix.yaml

@@ -20,6 +20,7 @@ jobs:
       sklearnDev:
         envFile: 'ci/environment-3.7.yaml'
         SKLEARN_DEV: "yes"
+        WRAPPERS: "no"
 
   steps:
   - bash: echo "##vso[task.prependpath]$CONDA/bin"
@@ -37,9 +38,20 @@ jobs:
       conda install -y -q pytorch cpuonly -c pytorch -n dask-ml-test
       source activate dask-ml-test
       pip install skorch
+      python -c "import torch; print('PyTorch ' + torch.__version__)"
+      python -c "import skorch; print('Skorch ' + skorch.__version__)"
     displayName: "install PyTorch"
     condition: eq(variables['Build.SourceBranch'], 'refs/heads/master')
 
+  - bash: |
+      source activate dask-ml-test
+      pip install tensorflow>=2.3.0
+      pip install scikeras>=0.1.8
+      python -c "import tensorflow as tf; print('TF ' + tf.__version__)"
+      python -c "import scikeras; print('SciKeras ' + scikeras.__version__)"
+    displayName: "install Tensorflow and SciKeras"
+    condition: eq(variables['Build.SourceBranch'], 'refs/heads/master')
+
   - script: |
       source activate dask-ml-test
       conda uninstall -y --force scikit-learn

docs/source/index.rst

@@ -133,6 +133,7 @@ Scikit-Learn should feel at home with Dask-ML.
    joblib.rst
    xgboost.rst
    pytorch.rst
+   keras.rst
 
 .. toctree::
    :maxdepth: 2

docs/source/keras.rst

@@ -0,0 +1,114 @@
+Keras and Tensorflow
+====================
+
+The package SciKeras_ brings a Scikit-learn API to Keras. This allows Dask-ML
+to be used seamlessly with Keras models.
+
+Installation
+------------
+
+Following the `Tensorflow install directions`_ and `SciKeras install guide`_,
+these packages need to be installed:
+
+.. code-block:: bash
+
+   $ pip install tensorflow>=2.3.0
+   $ pip install scikeras>=0.1.8
+
+These are the minimum versions that Dask-ML requires to use Tensorflow/Keras.
+
+.. _Tensorflow install directions: https://www.tensorflow.org/install
+.. _SciKeras install guide: https://github.com/adriangb/scikeras#installation
+
+Usage
+-----
+
+First, let's start by defining normal function to create our model. This is the
+normal way to create a `Keras Sequential model`_
+
+.. _Keras Sequential model: https://keras.io/api/models/sequential/
+
+.. code-block:: python
+
+   import tensorflow as tf
+   from tensorflow.keras.layers import Dense
+   from tensorflow.keras.models import Sequential
+
+   def build_model(lr=0.01, momentum=0.9):
+       layers = [Dense(512, input_shape=(784,), activation="relu"),
+                 Dense(10, input_shape=(512,), activation="softmax")]
+       model = Sequential(layers)
+
+       opt = tf.keras.optimizers.SGD(
+           learning_rate=lr, momentum=momentum, nesterov=True,
+       )
+       model.compile(loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"])
+       return model
+
+Now, we can use the SciKeras to create a Scikit-learn compatible model:
+
+.. code-block:: python
+
+   from scikeras.wrappers import KerasClassifier
+   niceties = dict(verbose=False)
+   model = KerasClassifier(build_fn=build_model, lr=0.1, momentum=0.9, **niceties)
+
+This model will work with all of Dask-ML: it can use NumPy arrays as inputs and
+obeys the Scikit-learn API. For example, it's possible to use Dask-ML to do the
+following:
+
+* Use Keras with Dask-ML's model selection, including
+  :class:`~dask_ml.model_selection.HyperbandSearchCV`.
+* Use Keras with Dask-ML's :class:`~dask_ml.wrappers.Incremental`.
+
+If we want to tune ``lr`` and ``momentum``, SciKeras requires that we pass
+``lr`` and ``momentum`` at initialization:
+
+.. code-block::
+
+   model = KerasClassifier(build_fn=build_model, lr=None, momentum=None, **niceties)
+
+.. _SciKeras: https://github.com/adriangb/scikeras
+
+SciKeras supports more model creation methods, including some that are
+backwards-compatible with Tensorflow. Refer to their documentation for details.
+
+Example: Hyperparameter Optimization
+------------------------------------
+
+If we wanted to, we could use the model above with
+:class:`~dask_ml.model_selection.HyperbandSearchCV`. Let's tune this model on
+the MNIST dataset:
+
+.. code-block:: python
+
+   from tensorflow.keras.datasets import mnist
+   from tensorflow.keras.utils import to_categorical
+   import numpy as np
+   from typing import Tuple
+
+   def get_mnist() -> Tuple[np.ndarray, np.ndarray]:
+       (X_train, y_train), _ = mnist.load_data()
+       X_train = X_train.reshape(X_train.shape[0], 784)
+       X_train = X_train.astype("float32")
+       X_train /= 255
+       return X_train, y_train
+
+And let's perform the basic task of tuning our SGD implementation:
+
+.. code-block:: python
+
+   from scipy.stats import loguniform, uniform
+   params = {"lr": loguniform(1e-3, 1e-1), "momentum": uniform(0, 1)}
+   X, y = get_mnist()
+
+Now, the search can be run:
+
+.. code-block:: python
+
+   from dask.distributed import Client
+   client = Client()
+
+   from dask_ml.model_selection import HyperbandSearchCV
+   search = HyperbandSearchCV(model, params, max_iter=27)
+   search.fit(X, y)

tests/model_selection/test_keras.py

@@ -0,0 +1,71 @@
+import numpy as np
+import pandas as pd
+import pytest
+from distributed import Nanny
+from distributed.utils_test import gen_cluster
+from packaging import version
+from scipy.stats import loguniform
+from sklearn.datasets import make_classification
+
+from dask_ml.model_selection import IncrementalSearchCV
+
+try:
+    import scikeras
+    import tensorflow as tf
+    from scikeras.wrappers import KerasClassifier
+    from tensorflow.keras.layers import Dense
+    from tensorflow.keras.models import Sequential
+
+    pytestmark = [
+        pytest.mark.skipif(
+            version.parse(tf.__version__) < version.parse("2.3.0"),
+            reason="pickle support",
+        ),
+        pytest.mark.skipif(
+            version.parse(scikeras.__version__) < version.parse("0.1.8"),
+            reason="partial_fit support",
+        ),
+    ]
+except ImportError:
+    pytestmark = pytest.mark.skip(reason="Missing tensorflow or scikeras")
+
+
+def _keras_build_fn(lr=0.01):
+    layers = [
+        Dense(512, input_shape=(784,), activation="relu"),
+        Dense(10, input_shape=(512,), activation="softmax"),
+    ]
+
+    model = Sequential(layers)
+
+    opt = tf.keras.optimizers.SGD(learning_rate=lr)
+    model.compile(loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"])
+    return model
+
+
+@gen_cluster(client=True, Worker=Nanny)
+def test_keras(c, s, a, b):
+    # Mirror the mnist dataset
+    X, y = make_classification(n_classes=10, n_features=784, n_informative=100)
+    X = X.astype("float32")
+    assert y.dtype == np.dtype("int64")
+
+    model = KerasClassifier(build_fn=_keras_build_fn, lr=0.01, verbose=False)
+    params = {"lr": loguniform(1e-3, 1e-1)}
+
+    search = IncrementalSearchCV(
+        model, params, max_iter=3, n_initial_parameters=5, decay_rate=None
+    )
+    yield search.fit(X, y)
+    #  search.fit(X, y)
+
+    assert search.best_score_ >= 0
+
+    # Make sure the model trains, and scores aren't constant
+    scores = {
+        ident: [h["score"] for h in hist]
+        for ident, hist in search.model_history_.items()
+    }
+    assert all(len(hist) == 3 for hist in scores.values())
+    nuniq_scores = [pd.Series(v).nunique() for v in scores.values()]
+    assert max(nuniq_scores) > 1