[tune] Remove keras dependency (#6827)

doc/examples/doc_code/tf_example.py

@@ -9,12 +9,12 @@ in the documentation.
 
 # yapf: disable
 # __tf_model_start__
-from tensorflow.keras import layers
 
 
 def create_keras_model():
-    import tensorflow as tf
-    model = tf.keras.Sequential()
+    from tensorflow import keras
+    from tensorflow.keras import layers
+    model = keras.Sequential()
     # Adds a densely-connected layer with 64 units to the model:
     model.add(layers.Dense(64, activation="relu", input_shape=(32, )))
     # Add another:
@@ -23,9 +23,9 @@ def create_keras_model():
     model.add(layers.Dense(10, activation="softmax"))
 
     model.compile(
-        optimizer=tf.keras.optimizers.RMSprop(0.01),
-        loss=tf.keras.losses.categorical_crossentropy,
-        metrics=[tf.keras.metrics.categorical_accuracy])
+        optimizer=keras.optimizers.RMSprop(0.01),
+        loss=keras.losses.categorical_crossentropy,
+        metrics=[keras.metrics.categorical_accuracy])
     return model
 # __tf_model_end__
 # yapf: enable

docker/examples/Dockerfile

@@ -9,7 +9,7 @@ RUN apt-get install -y zlib1g-dev libgl1-mesa-dev
 # The following is needed to support TensorFlow 1.14
 RUN conda remove -y --force wrapt
 RUN pip install -U pip
-RUN pip install gym[atari] opencv-python-headless tensorflow lz4 keras pytest-timeout smart_open tensorflow_probability
+RUN pip install gym[atari] opencv-python-headless tensorflow lz4 pytest-timeout smart_open tensorflow_probability
 RUN pip install -U h5py  # Mutes FutureWarnings
 RUN pip install --upgrade bayesian-optimization
 RUN pip install --upgrade hyperopt==0.1.2

python/ray/tune/examples/track_example.py

@@ -1,11 +1,10 @@
 import argparse
-import keras
-from keras.datasets import mnist
-from keras.models import Sequential
-from keras.layers import (Dense, Dropout, Flatten, Conv2D, MaxPooling2D)
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras.datasets import mnist
 
 from ray.tune import track
-from ray.tune.examples.utils import TuneReporterCallback, get_mnist_data
+from ray.tune.integration.keras import TuneReporterCallback
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
@@ -32,24 +31,18 @@ def train_mnist(args):
     batch_size = 128
     num_classes = 10
     epochs = 1 if args.smoke_test else 12
-    mnist.load()
-    x_train, y_train, x_test, y_test, input_shape = get_mnist_data()
 
-    model = Sequential()
-    model.add(
-        Conv2D(
-            32, kernel_size=(3, 3), activation="relu",
-            input_shape=input_shape))
-    model.add(Conv2D(64, (3, 3), activation="relu"))
-    model.add(MaxPooling2D(pool_size=(2, 2)))
-    model.add(Dropout(0.5))
-    model.add(Flatten())
-    model.add(Dense(args.hidden, activation="relu"))
-    model.add(Dropout(0.5))
-    model.add(Dense(num_classes, activation="softmax"))
+    (x_train, y_train), (x_test, y_test) = mnist.load_data()
+    x_train, x_test = x_train / 255.0, x_test / 255.0
+    model = tf.keras.models.Sequential([
+        tf.keras.layers.Flatten(input_shape=(28, 28)),
+        tf.keras.layers.Dense(args.hidden, activation="relu"),
+        tf.keras.layers.Dropout(0.2),
+        tf.keras.layers.Dense(num_classes, activation="softmax")
+    ])
 
     model.compile(
-        loss="categorical_crossentropy",
+        loss="sparse_categorical_crossentropy",
         optimizer=keras.optimizers.SGD(lr=args.lr, momentum=args.momentum),
         metrics=["accuracy"])
 
@@ -59,7 +52,7 @@ def train_mnist(args):
         batch_size=batch_size,
         epochs=epochs,
         validation_data=(x_test, y_test),
-        callbacks=[TuneReporterCallback(track.metric)])
+        callbacks=[TuneReporterCallback()])
     track.shutdown()
 
 

python/ray/tune/examples/tune_mnist_keras.py

@@ -3,7 +3,6 @@ import numpy as np
 from tensorflow.keras.datasets import mnist
 
 from ray.tune.integration.keras import TuneReporterCallback
-from ray.tune.examples.utils import get_mnist_data
 
 parser = argparse.ArgumentParser()
 parser.add_argument(
@@ -14,30 +13,21 @@ args, _ = parser.parse_known_args()
 def train_mnist(config, reporter):
     # https://github.com/tensorflow/tensorflow/issues/32159
     import tensorflow as tf
-    from tensorflow.keras.models import Sequential
-    from tensorflow.keras.layers import (Dense, Dropout, Flatten, Conv2D,
-                                         MaxPooling2D)
     batch_size = 128
     num_classes = 10
     epochs = 12
 
-    x_train, y_train, x_test, y_test, input_shape = get_mnist_data()
-
-    model = Sequential()
-    model.add(
-        Conv2D(
-            32, kernel_size=(3, 3), activation="relu",
-            input_shape=input_shape))
-    model.add(Conv2D(64, (3, 3), activation="relu"))
-    model.add(MaxPooling2D(pool_size=(2, 2)))
-    model.add(Dropout(0.5))
-    model.add(Flatten())
-    model.add(Dense(config["hidden"], activation="relu"))
-    model.add(Dropout(0.5))
-    model.add(Dense(num_classes, activation="softmax"))
+    (x_train, y_train), (x_test, y_test) = mnist.load_data()
+    x_train, x_test = x_train / 255.0, x_test / 255.0
+    model = tf.keras.models.Sequential([
+        tf.keras.layers.Flatten(input_shape=(28, 28)),
+        tf.keras.layers.Dense(config["hidden"], activation="relu"),
+        tf.keras.layers.Dropout(0.2),
+        tf.keras.layers.Dense(num_classes, activation="softmax")
+    ])
 
     model.compile(
-        loss=tf.keras.losses.categorical_crossentropy,
+        loss="sparse_categorical_crossentropy",
         optimizer=tf.keras.optimizers.SGD(
             lr=config["lr"], momentum=config["momentum"]),
         metrics=["accuracy"])

python/ray/tune/examples/utils.py

@@ -1,40 +1,9 @@
 import tensorflow as tf
-import keras
-from keras.datasets import mnist
-from keras import backend as K
 from sklearn.datasets import load_iris
 from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import OneHotEncoder
 
 
-def get_mnist_data():
-    img_rows, img_cols = 28, 28
-    num_classes = 10
-
-    # the data, split between train and test sets
-    (x_train, y_train), (x_test, y_test) = mnist.load_data()
-
-    if K.image_data_format() == "channels_first":
-        x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
-        x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
-        input_shape = (1, img_rows, img_cols)
-    else:
-        x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
-        x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
-        input_shape = (img_rows, img_cols, 1)
-
-    x_train = x_train.astype("float32")
-    x_test = x_test.astype("float32")
-    x_train /= 255
-    x_test /= 255
-
-    # convert class vectors to binary class matrices
-    y_train = keras.utils.to_categorical(y_train, num_classes)
-    y_test = keras.utils.to_categorical(y_test, num_classes)
-
-    return x_train, y_train, x_test, y_test, input_shape
-
-
 def get_iris_data(test_size=0.2):
     iris_data = load_iris()
     x = iris_data.data
@@ -48,11 +17,8 @@ def get_iris_data(test_size=0.2):
 def set_keras_threads(threads):
     # We set threads here to avoid contention, as Keras
     # is heavily parallelized across multiple cores.
-    K.set_session(
-        tf.Session(
-            config=tf.ConfigProto(
-                intra_op_parallelism_threads=threads,
-                inter_op_parallelism_threads=threads)))
+    tf.config.threading.set_inter_op_parallelism_threads(threads)
+    tf.config.threading.set_intra_op_parallelism_threads(threads)
 
 
 def TuneKerasCallback(*args, **kwargs):