chore: import upstream snapshot with attribution

python/ray/train/examples/tf/tensorflow_autoencoder_example.py

@@ -0,0 +1,177 @@
+# This example showcases how to use Tensorflow with Ray Train.
+# Original code:
+# https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+import ray
+from ray import train
+from ray.air.integrations.keras import ReportCheckpointCallback
+from ray.data.datasource import SimpleTensorFlowDatasource
+from ray.data.extensions import TensorArray
+from ray.train import Result, ScalingConfig
+from ray.train.tensorflow import TensorflowTrainer, prepare_dataset_shard
+
+
+def get_dataset(split_type="train"):
+    def dataset_factory():
+        return tfds.load("mnist", split=[split_type], as_supervised=True)[0].take(128)
+
+    dataset = ray.data.read_datasource(
+        SimpleTensorFlowDatasource(), dataset_factory=dataset_factory
+    )
+
+    def normalize_images(x):
+        x = np.float32(x.numpy()) / 255.0
+        x = np.reshape(x, (-1,))
+        return x
+
+    def preprocess_dataset(batch):
+        return [
+            (normalize_images(image), normalize_images(image)) for image, _ in batch
+        ]
+
+    dataset = dataset.map_batches(preprocess_dataset)
+
+    def convert_batch_to_pandas(batch):
+
+        images = [TensorArray(image) for image, _ in batch]
+        # because we did autoencoder here
+        df = pd.DataFrame({"image": images, "label": images})
+        return df
+
+    dataset = dataset.map_batches(convert_batch_to_pandas)
+    return dataset
+
+
+def build_autoencoder_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=(784,)),
+            # encoder
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(64, activation="relu"),
+            tf.keras.layers.Dense(32, activation="relu"),
+            # decoder
+            tf.keras.layers.Dense(64, activation="relu"),
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(784, activation="sigmoid"),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+
+    per_worker_batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+
+    dataset_shard = train.get_dataset_shard("train")
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_autoencoder_model()
+        learning_rate = config.get("lr", 0.001)
+        multi_worker_model.compile(
+            loss=tf.keras.losses.BinaryCrossentropy(),
+            optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
+            metrics=[
+                "binary_crossentropy",
+            ],
+        )
+
+    def to_tf_dataset(dataset, batch_size):
+        def to_tensor_iterator():
+            for batch in dataset.iter_tf_batches(
+                batch_size=batch_size, dtypes=tf.float32
+            ):
+                yield batch["image"], batch["label"]
+
+        output_signature = (
+            tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
+            tf.TensorSpec(shape=(None, 784), dtype=tf.float32),
+        )
+        tf_dataset = tf.data.Dataset.from_generator(
+            to_tensor_iterator, output_signature=output_signature
+        )
+        return prepare_dataset_shard(tf_dataset)
+
+    results = []
+    for epoch in range(epochs):
+        tf_dataset = to_tf_dataset(
+            dataset=dataset_shard,
+            batch_size=per_worker_batch_size,
+        )
+        history = multi_worker_model.fit(
+            tf_dataset, callbacks=[ReportCheckpointCallback()]
+        )
+        results.append(history.history)
+    return results
+
+
+def train_tensorflow_mnist(
+    num_workers: int = 2, use_gpu: bool = False, epochs: int = 4
+) -> Result:
+    train_dataset = get_dataset(split_type="train")
+    config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        datasets={"train": train_dataset},
+        scaling_config=scaling_config,
+    )
+
+    results = trainer.fit()
+    print(results.metrics)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        result = train_tensorflow_mnist(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        result = train_tensorflow_mnist(
+            num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
+        )
+    print(result)

python/ray/train/examples/tf/tensorflow_mnist_example.py

@@ -0,0 +1,138 @@
+# This example showcases how to use Tensorflow with Ray Train.
+# Original code:
+# https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+import json
+
+import numpy as np
+import tensorflow as tf
+from filelock import FileLock
+
+from ray.air.integrations.keras import ReportCheckpointCallback
+from ray.train import Result, RunConfig, ScalingConfig
+from ray.train.tensorflow import TensorflowTrainer
+
+
+def mnist_dataset(batch_size: int) -> tf.data.Dataset:
+    with FileLock(os.path.expanduser("~/.mnist_lock")):
+        (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
+    # The `x` arrays are in uint8 and have values in the [0, 255] range.
+    # You need to convert them to float32 with values in the [0, 1] range.
+    x_train = x_train / np.float32(255)
+    y_train = y_train.astype(np.int64)
+    train_dataset = (
+        tf.data.Dataset.from_tensor_slices((x_train, y_train))
+        .shuffle(60000)
+        .repeat()
+        .batch(batch_size)
+    )
+    return train_dataset
+
+
+def build_cnn_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.Input(shape=(28, 28)),
+            tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
+            tf.keras.layers.Conv2D(32, 3, activation="relu"),
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(10),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    per_worker_batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+    steps_per_epoch = config.get("steps_per_epoch", 70)
+
+    tf_config = json.loads(os.environ["TF_CONFIG"])
+    num_workers = len(tf_config["cluster"]["worker"])
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    global_batch_size = per_worker_batch_size * num_workers
+    multi_worker_dataset = mnist_dataset(global_batch_size)
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_cnn_model()
+        learning_rate = config.get("lr", 0.001)
+        multi_worker_model.compile(
+            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate),
+            metrics=["accuracy"],
+        )
+
+    history = multi_worker_model.fit(
+        multi_worker_dataset,
+        epochs=epochs,
+        steps_per_epoch=steps_per_epoch,
+        callbacks=[ReportCheckpointCallback()],
+    )
+    results = history.history
+    return results
+
+
+def train_tensorflow_mnist(
+    num_workers: int = 2,
+    use_gpu: bool = False,
+    epochs: int = 4,
+    storage_path: str = None,
+) -> Result:
+    config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
+        run_config=RunConfig(storage_path=storage_path),
+    )
+    results = trainer.fit()
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=3, help="Number of epochs to train for."
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    import ray
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        train_tensorflow_mnist(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        train_tensorflow_mnist(
+            num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
+        )

python/ray/train/examples/tf/tensorflow_quick_start.py

@@ -0,0 +1,90 @@
+# ruff: noqa
+# fmt: off
+# isort: skip_file
+
+# __tf_setup_begin__
+import os
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import sys
+import numpy as np
+
+if sys.version_info >= (3, 12):
+    # Tensorflow is not installed for Python 3.12 because of keras compatibility.
+    sys.exit(0)
+else:
+    import tensorflow as tf
+
+def mnist_dataset(batch_size):
+    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
+    # The `x` arrays are in uint8 and have values in the [0, 255] range.
+    # You need to convert them to float32 with values in the [0, 1] range.
+    x_train = x_train / np.float32(255)
+    y_train = y_train.astype(np.int64)
+    train_dataset = tf.data.Dataset.from_tensor_slices(
+        (x_train, y_train)).shuffle(60000).repeat().batch(batch_size)
+    return train_dataset
+
+
+def build_and_compile_cnn_model():
+    model = tf.keras.Sequential([
+        tf.keras.layers.InputLayer(input_shape=(28, 28)),
+        tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
+        tf.keras.layers.Conv2D(32, 3, activation='relu'),
+        tf.keras.layers.Flatten(),
+        tf.keras.layers.Dense(128, activation='relu'),
+        tf.keras.layers.Dense(10)
+    ])
+    model.compile(
+        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
+        metrics=['accuracy'])
+    return model
+# __tf_setup_end__
+
+# __tf_single_begin__
+def train_func():
+    batch_size = 64
+    single_worker_dataset = mnist_dataset(batch_size)
+    single_worker_model = build_and_compile_cnn_model()
+    single_worker_model.fit(single_worker_dataset, epochs=3, steps_per_epoch=70)
+# __tf_single_end__
+
+# __tf_distributed_begin__
+import json
+import os
+
+def train_func_distributed():
+    per_worker_batch_size = 64
+    # This environment variable will be set by Ray Train.
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    num_workers = len(tf_config['cluster']['worker'])
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+
+    global_batch_size = per_worker_batch_size * num_workers
+    multi_worker_dataset = mnist_dataset(global_batch_size)
+
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_and_compile_cnn_model()
+
+    multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70)
+# __tf_distributed_end__
+
+if __name__ == "__main__":
+    # __tf_single_run_begin__
+    train_func()
+    # __tf_single_run_end__
+
+    # __tf_trainer_begin__
+    from ray.train.tensorflow import TensorflowTrainer
+    from ray.train import ScalingConfig
+
+    # For GPU Training, set `use_gpu` to True.
+    use_gpu = False
+
+    trainer = TensorflowTrainer(train_func_distributed, scaling_config=ScalingConfig(num_workers=4, use_gpu=use_gpu))
+
+    trainer.fit()
+    # __tf_trainer_end__

python/ray/train/examples/tf/tensorflow_regression_example.py

@@ -0,0 +1,115 @@
+import os
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+import argparse
+import sys
+
+import ray
+from ray import train
+from ray.data.preprocessors import Concatenator
+from ray.train import Result, ScalingConfig
+
+if sys.version_info >= (3, 12):
+    # Skip this test in Python 3.12+ because TensorFlow is not supported.
+    sys.exit(0)
+else:
+    import tensorflow as tf
+
+    from ray.train.tensorflow import TensorflowTrainer
+    from ray.train.tensorflow.keras import ReportCheckpointCallback
+
+
+def build_model() -> tf.keras.Model:
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(input_shape=(100,)),
+            tf.keras.layers.Dense(10),
+            tf.keras.layers.Dense(1),
+        ]
+    )
+    return model
+
+
+def train_func(config: dict):
+    batch_size = config.get("batch_size", 64)
+    epochs = config.get("epochs", 3)
+
+    strategy = tf.distribute.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+        # Model building/compiling need to be within `strategy.scope()`.
+        multi_worker_model = build_model()
+        multi_worker_model.compile(
+            optimizer=tf.keras.optimizers.SGD(learning_rate=config.get("lr", 1e-3)),
+            loss=tf.keras.losses.mean_absolute_error,
+            metrics=[tf.keras.metrics.mean_squared_error],
+        )
+
+    dataset = train.get_dataset_shard("train")
+
+    results = []
+    for _ in range(epochs):
+        tf_dataset = dataset.to_tf(
+            feature_columns="x", label_columns="y", batch_size=batch_size
+        )
+        history = multi_worker_model.fit(
+            tf_dataset, callbacks=[ReportCheckpointCallback()]
+        )
+        results.append(history.history)
+    return results
+
+
+def train_tensorflow_regression(num_workers: int = 2, use_gpu: bool = False) -> Result:
+    dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv")
+    columns_to_concatenate = [f"x{i:03}" for i in range(100)]
+    preprocessor = Concatenator(columns=columns_to_concatenate, output_column_name="x")
+    dataset = preprocessor.fit_transform(dataset)
+
+    config = {"lr": 1e-3, "batch_size": 32, "epochs": 4}
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
+    trainer = TensorflowTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=config,
+        scaling_config=scaling_config,
+        datasets={"train": dataset},
+    )
+    results = trainer.fit()
+    print(results.metrics)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--address", required=False, type=str, help="the address to use for Ray"
+    )
+    parser.add_argument(
+        "--num-workers",
+        "-n",
+        type=int,
+        default=2,
+        help="Sets number of workers for training.",
+    )
+    parser.add_argument(
+        "--use-gpu", action="store_true", default=False, help="Enables GPU training"
+    )
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    if args.smoke_test:
+        # 2 workers, 1 for trainer, 1 for datasets
+        num_gpus = args.num_workers if args.use_gpu else 0
+        ray.init(num_cpus=4, num_gpus=num_gpus)
+        result = train_tensorflow_regression(num_workers=2, use_gpu=args.use_gpu)
+    else:
+        ray.init(address=args.address)
+        result = train_tensorflow_regression(
+            num_workers=args.num_workers, use_gpu=args.use_gpu
+        )
+    print(result)