wehub-resource-sync
139 lines
4.2 KiB
Python
139 lines
4.2 KiB
Python
# 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
|
|
)
|