[training] Tensorflow interface for MultiNode SGD (#5440)

python/ray/experimental/sgd/examples/cifar_tf_example.py

@@ -0,0 +1,222 @@
+"""
+#Train a simple deep CNN on the CIFAR10 small images dataset.
+
+It gets to 75% validation accuracy in 25 epochs, and 79% after 50 epochs.
+(it"s still underfitting at that point, though).
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import tensorflow as tf
+from tensorflow import keras
+
+from tensorflow.keras.datasets import cifar10
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense, Dropout, Activation, Flatten
+from tensorflow.keras.layers import Conv2D, MaxPooling2D
+import os
+from filelock import FileLock
+
+import ray
+from ray.experimental.sgd.tf.tf_trainer import TFTrainer
+
+num_classes = 10
+
+
+def fetch_keras_data():
+    # The data, split between train and test sets:
+    with FileLock(os.path.expanduser("~/.cifar.lock")):
+        (x_train, y_train), (x_test, y_test) = cifar10.load_data()
+
+    # 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)
+
+    x_train = x_train.astype("float32")
+    x_test = x_test.astype("float32")
+    x_train /= 255
+    x_test /= 255
+    return (x_train, y_train), (x_test, y_test)
+
+
+(x_train, y_train), (x_test, y_test) = fetch_keras_data()
+input_shape = x_train.shape[1:]
+
+
+def create_model(config):
+    model = Sequential()
+    model.add(Conv2D(32, (3, 3), padding="same", input_shape=input_shape))
+    model.add(Activation("relu"))
+    model.add(Conv2D(32, (3, 3)))
+    model.add(Activation("relu"))
+    model.add(MaxPooling2D(pool_size=(2, 2)))
+    model.add(Dropout(0.25))
+
+    model.add(Conv2D(64, (3, 3), padding="same"))
+    model.add(Activation("relu"))
+    model.add(Conv2D(64, (3, 3)))
+    model.add(Activation("relu"))
+    model.add(MaxPooling2D(pool_size=(2, 2)))
+    model.add(Dropout(0.25))
+
+    model.add(Flatten())
+    model.add(Dense(512))
+    model.add(Activation("relu"))
+    model.add(Dropout(0.5))
+    model.add(Dense(num_classes))
+    model.add(Activation("softmax"))
+
+    # initiate RMSprop optimizer
+    opt = keras.optimizers.RMSprop(lr=0.001, decay=1e-6)
+
+    # Let"s train the model using RMSprop
+    model.compile(
+        loss="categorical_crossentropy", optimizer=opt, metrics=["accuracy"])
+    return model
+
+
+def data_creator(config):
+    batch_size = config["batch_size"]
+    (x_train, y_train), (x_test, y_test) = fetch_keras_data()
+    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+
+    # Repeat is needed to avoid
+    train_dataset = train_dataset.repeat().shuffle(
+        len(x_train)).batch(batch_size)
+    test_dataset = test_dataset.repeat().batch(batch_size)
+    return train_dataset, test_dataset
+
+
+def _make_generator(x_train, y_train, batch_size):
+    # This will do preprocessing and realtime data augmentation:
+    datagen = ImageDataGenerator(
+        featurewise_center=False,  # set input mean to 0 over the dataset
+        samplewise_center=False,  # set each sample mean to 0
+        # divide inputs by std of the dataset
+        featurewise_std_normalization=False,
+        samplewise_std_normalization=False,  # divide each input by its std
+        zca_whitening=False,  # apply ZCA whitening
+        zca_epsilon=1e-06,  # epsilon for ZCA whitening
+        # randomly rotate images in the range (degrees, 0 to 180)
+        rotation_range=0,
+        # randomly shift images horizontally (fraction of total width)
+        width_shift_range=0.1,
+        # randomly shift images vertically (fraction of total height)
+        height_shift_range=0.1,
+        shear_range=0.,  # set range for random shear
+        zoom_range=0.,  # set range for random zoom
+        channel_shift_range=0.,  # set range for random channel shifts
+        # set mode for filling points outside the input boundaries
+        fill_mode="nearest",
+        cval=0.,  # value used for fill_mode = "constant"
+        horizontal_flip=True,  # randomly flip images
+        vertical_flip=False,  # randomly flip images
+        # set rescaling factor (applied before any other transformation)
+        rescale=None,
+        # set function that will be applied on each input
+        preprocessing_function=None,
+        # image data format, either "channels_first" or "channels_last"
+        data_format=None,
+        # fraction of images reserved for validation (strictly between 0 and 1)
+        validation_split=0.0)
+
+    # Compute quantities required for feature-wise normalization
+    # (std, mean, and principal components if ZCA whitening is applied).
+    datagen.fit(x_train)
+    return datagen.flow(x_train, y_train, batch_size=batch_size)
+
+
+def data_augmentation_creator(config):
+    batch_size = config["batch_size"]
+    (x_train, y_train), (x_test, y_test) = fetch_keras_data()
+    trainset = tf.data.Dataset.from_generator(
+        lambda: _make_generator(x_train, y_train, batch_size),
+        output_types=(tf.float32, tf.float32),
+        # https://github.com/tensorflow/tensorflow/issues/24520
+        output_shapes=(tf.TensorShape((None, None, None, None)),
+                       tf.TensorShape((None, 10))))
+    trainset = trainset.repeat()
+
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+    test_dataset = test_dataset.repeat().batch(batch_size)
+    return trainset, test_dataset
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--redis-address",
+        required=False,
+        type=str,
+        help="the address to use for Redis")
+    parser.add_argument(
+        "--num-replicas",
+        "-n",
+        type=int,
+        default=1,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=512,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Enables GPU training")
+    parser.add_argument(
+        "--augment-data",
+        action="store_true",
+        default=False,
+        help="Sets data augmentation.")
+    parser.add_argument(
+        "--smoke-test",
+        action="store_true",
+        default=False,
+        help="Finish quickly for testing. Assume False for users.")
+
+    args, _ = parser.parse_known_args()
+    ray.init(redis_address=args.redis_address)
+    data_size = 60000
+    test_size = 10000
+    batch_size = args.batch_size
+
+    num_train_steps = 10 if args.smoke_test else data_size // batch_size
+    num_eval_steps = 10 if args.smoke_test else test_size // batch_size
+
+    trainer = TFTrainer(
+        model_creator=create_model,
+        data_creator=(data_augmentation_creator
+                      if args.augment_data else data_creator),
+        num_replicas=args.num_replicas,
+        use_gpu=args.use_gpu,
+        verbose=True,
+        config={
+            "batch_size": batch_size,
+            "fit_config": {
+                "steps_per_epoch": num_train_steps,
+            },
+            "evaluate_config": {
+                "steps": num_eval_steps,
+            }
+        })
+
+    for i in range(3):
+        # Trains num epochs
+        train_stats1 = trainer.train()
+        train_stats1.update(trainer.validate())
+        print("iter {}:".format(i), train_stats1)
+
+    model = trainer.get_model()
+    trainer.shutdown()
+    dataset, test_dataset = data_augmentation_creator(
+        dict(batch_size=batch_size))
+    model.fit(dataset, steps_per_epoch=num_train_steps, epochs=1)
+    scores = model.evaluate(test_dataset, steps=num_eval_steps)
+    print("Test loss:", scores[0])
+    print("Test accuracy:", scores[1])

python/ray/experimental/sgd/examples/tensorflow_train_example.py

@@ -0,0 +1,135 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+from tensorflow.data import Dataset
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense
+import numpy as np
+
+import ray
+from ray import tune
+from ray.experimental.sgd.tf.tf_trainer import TFTrainer, TFTrainable
+
+NUM_TRAIN_SAMPLES = 1000
+NUM_TEST_SAMPLES = 400
+
+
+def create_config(batch_size):
+    return {
+        "batch_size": batch_size,
+        "fit_config": {
+            "steps_per_epoch": NUM_TRAIN_SAMPLES // batch_size
+        },
+        "evaluate_config": {
+            "steps": NUM_TEST_SAMPLES // batch_size,
+        }
+    }
+
+
+def linear_dataset(a=2, size=1000):
+    x = np.random.rand(size)
+    y = x / 2
+
+    x = x.reshape((-1, 1))
+    y = y.reshape((-1, 1))
+
+    return x, y
+
+
+def simple_dataset(config):
+    batch_size = config["batch_size"]
+    x_train, y_train = linear_dataset(size=NUM_TRAIN_SAMPLES)
+    x_test, y_test = linear_dataset(size=NUM_TEST_SAMPLES)
+
+    train_dataset = Dataset.from_tensor_slices((x_train, y_train))
+    test_dataset = Dataset.from_tensor_slices((x_test, y_test))
+    train_dataset = train_dataset.shuffle(NUM_TRAIN_SAMPLES).repeat().batch(
+        batch_size)
+    test_dataset = test_dataset.repeat().batch(batch_size)
+
+    return train_dataset, test_dataset
+
+
+def simple_model(config):
+    model = Sequential([Dense(10, input_shape=(1, )), Dense(1)])
+
+    model.compile(
+        optimizer="sgd",
+        loss="mean_squared_error",
+        metrics=["mean_squared_error"])
+
+    return model
+
+
+def train_example(num_replicas=1, batch_size=128, use_gpu=False):
+    trainer = TFTrainer(
+        model_creator=simple_model,
+        data_creator=simple_dataset,
+        num_replicas=num_replicas,
+        use_gpu=use_gpu,
+        verbose=True,
+        config=create_config(batch_size))
+
+    train_stats1 = trainer.train()
+    train_stats1.update(trainer.validate())
+    print(train_stats1)
+
+    train_stats2 = trainer.train()
+    train_stats2.update(trainer.validate())
+    print(train_stats2)
+
+    val_stats = trainer.validate()
+    print(val_stats)
+    print("success!")
+
+
+def tune_example(num_replicas=1, use_gpu=False):
+    config = {
+        "model_creator": tune.function(simple_model),
+        "data_creator": tune.function(simple_dataset),
+        "num_replicas": num_replicas,
+        "use_gpu": use_gpu,
+        "trainer_config": create_config(batch_size=128)
+    }
+
+    analysis = tune.run(
+        TFTrainable,
+        num_samples=2,
+        config=config,
+        stop={"training_iteration": 2},
+        verbose=1)
+
+    return analysis.get_best_config(metric="validation_loss", mode="min")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--redis-address",
+        required=False,
+        type=str,
+        help="the address to use for Redis")
+    parser.add_argument(
+        "--num-replicas",
+        "-n",
+        type=int,
+        default=1,
+        help="Sets number of replicas for training.")
+    parser.add_argument(
+        "--use-gpu",
+        action="store_true",
+        default=False,
+        help="Enables GPU training")
+    parser.add_argument(
+        "--tune", action="store_true", default=False, help="Tune training")
+
+    args, _ = parser.parse_known_args()
+
+    ray.init(redis_address=args.redis_address)
+
+    if args.tune:
+        tune_example(num_replicas=args.num_replicas, use_gpu=args.use_gpu)
+    else:
+        train_example(num_replicas=args.num_replicas, use_gpu=args.use_gpu)

python/ray/experimental/sgd/examples/tf-example-sgd.yaml

@@ -0,0 +1,73 @@
+# An unique identifier for the head node and workers of this cluster.
+cluster_name: sgd-tf
+
+# The maximum number of workers nodes to launch in addition to the head
+# node. This takes precedence over min_workers. min_workers default to 0.
+min_workers: 3
+initial_workers: 3
+max_workers: 3
+
+target_utilization_fraction: 0.9
+
+# If a node is idle for this many minutes, it will be removed.
+idle_timeout_minutes: 20
+# docker:
+#     image: tensorflow/tensorflow:1.5.0-py3
+#     container_name: ray_docker
+
+# Cloud-provider specific configuration.
+provider:
+    type: aws
+    region: us-east-1
+    availability_zone: us-east-1e
+
+# How Ray will authenticate with newly launched nodes.
+auth:
+    ssh_user: ubuntu
+
+head_node:
+    InstanceType: g3.8xlarge
+    ImageId: ami-0757fc5a639fe7666
+    #    InstanceMarketOptions:
+    #        MarketType: spot
+    #        SpotOptions:
+    #            MaxPrice: "9.0"
+
+
+worker_nodes:
+    InstanceType: g3.8xlarge
+    ImageId: ami-0757fc5a639fe7666
+    # InstanceMarketOptions:
+    #     MarketType: spot
+        # SpotOptions:
+        #     MaxPrice: "9.0"
+
+    #     # Run workers on spot by default. Comment this out to use on-demand.
+    #     InstanceMarketOptions:
+    #         MarketType: spot
+
+setup_commands:
+    - conda install setuptools=41.0.1=py36_0 wrapt=1.11.2 --yes # workaround to fix wrapt error
+    - ray || pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-0.8.0.dev3-cp36-cp36m-manylinux1_x86_64.whl
+    - pip install -U ipdb ray[rllib]
+    - pip install tensorflow==2.0.0-rc0
+
+file_mounts: {
+}
+
+# Custom commands that will be run on the head node after common setup.
+head_setup_commands: []
+
+# Custom commands that will be run on worker nodes after common setup.
+worker_setup_commands: []
+
+# # Command to start ray on the head node. You don't need to change this.
+head_start_ray_commands:
+    - ray stop
+    - ray start --head --redis-port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml --object-store-memory=1000000000
+
+# Command to start ray on worker nodes. You don't need to change this.
+worker_start_ray_commands:
+    - ray stop
+    - ray start --redis-address=$RAY_HEAD_IP:6379 --object-manager-port=8076 --object-store-memory=1000000000
+

python/ray/experimental/sgd/pytorch/pytorch_runner.py

@@ -7,7 +7,8 @@ import torch
 import torch.utils.data
 
 import ray
-from ray.experimental.sgd.pytorch import utils
+from ray.experimental.sgd.pytorch import pytorch_utils
+from ray.experimental.sgd import utils
 
 logger = logging.getLogger(__name__)
 
@@ -89,8 +90,8 @@ class PyTorchRunner(object):
         """Runs a training epoch and updates the model parameters."""
         logger.debug("Begin Training Epoch {}".format(self.epoch + 1))
         with self._timers["training"]:
-            train_stats = utils.train(self.train_loader, self.model,
-                                      self.criterion, self.optimizer)
+            train_stats = pytorch_utils.train(self.train_loader, self.model,
+                                              self.criterion, self.optimizer)
             train_stats["epoch"] = self.epoch
 
         self.epoch += 1
@@ -101,8 +102,8 @@ class PyTorchRunner(object):
     def validate(self):
         """Evaluates the model on the validation data set."""
         with self._timers["validation"]:
-            validation_stats = utils.validate(self.validation_loader,
-                                              self.model, self.criterion)
+            validation_stats = pytorch_utils.validate(
+                self.validation_loader, self.model, self.criterion)
 
         validation_stats.update(self.stats())
         return validation_stats

python/ray/experimental/sgd/pytorch/pytorch_trainer.py

@@ -15,7 +15,8 @@ from ray.tune.resources import Resources
 from ray.experimental.sgd.pytorch.pytorch_runner import PyTorchRunner
 from ray.experimental.sgd.pytorch.distributed_pytorch_runner import (
     DistributedPyTorchRunner)
-from ray.experimental.sgd.pytorch import utils
+from ray.experimental.sgd.pytorch import pytorch_utils
+from ray.experimental.sgd import utils
 
 logger = logging.getLogger(__name__)
 
@@ -30,7 +31,7 @@ class PyTorchTrainer(object):
     def __init__(self,
                  model_creator,
                  data_creator,
-                 optimizer_creator=utils.sgd_mse_optimizer,
+                 optimizer_creator=pytorch_utils.sgd_mse_optimizer,
                  config=None,
                  num_replicas=1,
                  use_gpu=False,

python/ray/experimental/sgd/pytorch/pytorch_utils.py

@@ -0,0 +1,107 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+import torch
+import torch.nn as nn
+
+from ray.experimental.sgd import utils
+
+
+def train(train_iterator, model, criterion, optimizer):
+    """Runs 1 training epoch"""
+    batch_time = utils.AverageMeter()
+    data_time = utils.AverageMeter()
+    losses = utils.AverageMeter()
+
+    timers = {k: utils.TimerStat() for k in ["d2h", "fwd", "grad", "apply"]}
+
+    # switch to train mode
+    model.train()
+
+    end = time.time()
+
+    for i, (features, target) in enumerate(train_iterator):
+        # measure data loading time
+        data_time.update(time.time() - end)
+
+        # Create non_blocking tensors for distributed training
+        with timers["d2h"]:
+            if torch.cuda.is_available():
+                features = features.cuda(non_blocking=True)
+                target = target.cuda(non_blocking=True)
+
+        # compute output
+        with timers["fwd"]:
+            output = model(features)
+            loss = criterion(output, target)
+
+            # measure accuracy and record loss
+            losses.update(loss.item(), features.size(0))
+
+        with timers["grad"]:
+            # compute gradients in a backward pass
+            optimizer.zero_grad()
+            loss.backward()
+
+        with timers["apply"]:
+            # Call step of optimizer to update model params
+            optimizer.step()
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+    stats = {
+        "batch_time": batch_time.avg,
+        "batch_processed": losses.count,
+        "train_loss": losses.avg,
+        "data_time": data_time.avg,
+    }
+    stats.update({k: t.mean for k, t in timers.items()})
+    return stats
+
+
+def validate(val_loader, model, criterion):
+    batch_time = utils.AverageMeter()
+    losses = utils.AverageMeter()
+
+    # switch to evaluate mode
+    model.eval()
+
+    with torch.no_grad():
+        end = time.time()
+        for i, (features, target) in enumerate(val_loader):
+
+            if torch.cuda.is_available():
+                features = features.cuda(non_blocking=True)
+                target = target.cuda(non_blocking=True)
+
+            # compute output
+            output = model(features)
+            loss = criterion(output, target)
+
+            # measure accuracy and record loss
+            losses.update(loss.item(), features.size(0))
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+    stats = {"batch_time": batch_time.avg, "validation_loss": losses.avg}
+    return stats
+
+
+def sgd_mse_optimizer(model, config):
+    """Returns the mean squared error criterion and SGD optimizer.
+
+    Args:
+        model (torch.nn.Module): the model to optimize.
+        config (dict): configuration for the optimizer.
+            lr (float): the learning rate. defaults to 0.01.
+    """
+    learning_rate = config.get("lr", 0.01)
+    criterion = nn.MSELoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
+    return criterion, optimizer

python/ray/experimental/sgd/tests/pytorch_utils.py

@@ -1,40 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.utils.data
-
-
-class LinearDataset(torch.utils.data.Dataset):
-    """y = a * x + b"""
-
-    def __init__(self, a, b, size=1000):
-        x = np.random.random(size).astype(np.float32) * 10
-        x = np.arange(0, 10, 10 / size, dtype=np.float32)
-        self.x = torch.from_numpy(x)
-        self.y = torch.from_numpy(a * x + b)
-
-    def __getitem__(self, index):
-        return self.x[index, None], self.y[index, None]
-
-    def __len__(self):
-        return len(self.x)
-
-
-def model_creator(config):
-    return nn.Linear(1, 1)
-
-
-def optimizer_creator(model, config):
-    """Returns criterion, optimizer"""
-    criterion = nn.MSELoss()
-    optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
-    return criterion, optimizer
-
-
-def data_creator(config):
-    """Returns training set, validation set"""
-    return LinearDataset(2, 5), LinearDataset(2, 5, size=400)

python/ray/experimental/sgd/tests/test_pytorch.py

@@ -12,7 +12,7 @@ from ray import tune
 from ray.tests.conftest import ray_start_2_cpus  # noqa: F401
 from ray.experimental.sgd.pytorch import PyTorchTrainer, PyTorchTrainable
 
-from ray.experimental.sgd.tests.pytorch_utils import (
+from ray.experimental.sgd.examples.train_example import (
     model_creator, optimizer_creator, data_creator)
 
 

python/ray/experimental/sgd/tests/test_tensorflow.py

@@ -0,0 +1,135 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import pytest
+import tempfile
+import numpy as np
+import shutil
+
+from ray import tune
+from ray.tests.conftest import ray_start_2_cpus  # noqa: F401
+from ray.experimental.sgd.tf import TFTrainer, TFTrainable
+
+from ray.experimental.sgd.examples.tensorflow_train_example import (
+    simple_model, simple_dataset)
+
+SIMPLE_CONFIG = {
+    "batch_size": 128,
+    "fit_config": {
+        "steps_per_epoch": 3,
+    },
+    "evaluate_config": {
+        "steps": 3,
+    }
+}
+
+
+@pytest.mark.parametrize(  # noqa: F811
+    "num_replicas", [1, 2])
+def test_train(ray_start_2_cpus, num_replicas):  # noqa: F811
+    trainer = TFTrainer(
+        model_creator=simple_model,
+        data_creator=simple_dataset,
+        num_replicas=num_replicas,
+        config=SIMPLE_CONFIG)
+
+    train_stats1 = trainer.train()
+    train_stats1.update(trainer.validate())
+
+    train_stats2 = trainer.train()
+    train_stats2.update(trainer.validate())
+
+
+@pytest.mark.parametrize(  # noqa: F811
+    "num_replicas", [1, 2])
+def test_tune_train(ray_start_2_cpus, num_replicas):  # noqa: F811
+
+    config = {
+        "model_creator": tune.function(simple_model),
+        "data_creator": tune.function(simple_dataset),
+        "num_replicas": num_replicas,
+        "use_gpu": False,
+        "trainer_config": SIMPLE_CONFIG
+    }
+
+    tune.run(
+        TFTrainable,
+        num_samples=2,
+        config=config,
+        stop={"training_iteration": 2},
+        verbose=1)
+
+
+@pytest.mark.parametrize(  # noqa: F811
+    "num_replicas", [1, 2])
+def test_save_and_restore(ray_start_2_cpus, num_replicas):  # noqa: F811
+    trainer1 = TFTrainer(
+        model_creator=simple_model,
+        data_creator=simple_dataset,
+        num_replicas=num_replicas,
+        config=SIMPLE_CONFIG)
+    trainer1.train()
+
+    tmpdir = tempfile.mkdtemp()
+    filename = os.path.join(tmpdir, "checkpoint")
+    trainer1.save(filename)
+
+    model1 = trainer1.get_model()
+    trainer1.shutdown()
+
+    trainer2 = TFTrainer(
+        model_creator=simple_model,
+        data_creator=simple_dataset,
+        num_replicas=num_replicas,
+        config=SIMPLE_CONFIG)
+    trainer2.restore(filename)
+
+    model2 = trainer2.get_model()
+    trainer2.shutdown()
+
+    shutil.rmtree(tmpdir)
+
+    model1_config = model1.get_config()
+    model2_config = model2.get_config()
+    assert _compare(model1_config, model2_config, skip_keys=["name"])
+
+    model1_weights = model1.get_weights()
+    model2_weights = model2.get_weights()
+    assert _compare(model1_weights, model2_weights)
+
+    model1_opt_weights = model1.optimizer.get_weights()
+    model2_opt_weights = model2.optimizer.get_weights()
+    assert _compare(model1_opt_weights, model2_opt_weights)
+
+
+def _compare(d1, d2, skip_keys=None):
+    """Compare two lists or dictionaries or array"""
+    if type(d1) != type(d2):
+        return False
+
+    if isinstance(d1, dict):
+        if set(d1) != set(d2):
+            return False
+
+        for key in d1:
+            if skip_keys is not None and key in skip_keys:
+                continue
+
+            if not _compare(d1[key], d2[key], skip_keys=skip_keys):
+                return False
+
+    elif isinstance(d1, list):
+        for i, _ in enumerate(d1):
+            if not _compare(d1[i], d2[i], skip_keys=skip_keys):
+                return False
+
+    elif isinstance(d1, np.ndarray):
+        if not np.array_equal(d1, d2):
+            return False
+    else:
+        if d1 != d2:
+            return False
+
+    return True

python/ray/experimental/sgd/tf/__init__.py

@@ -0,0 +1,7 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from ray.experimental.sgd.tf.tf_trainer import (TFTrainer, TFTrainable)
+
+__all__ = ["TFTrainer", "TFTrainable"]

python/ray/experimental/sgd/tf/tf_runner.py

@@ -0,0 +1,157 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+import json
+import os
+import numpy as np
+
+import ray
+import ray.services
+from ray.experimental.sgd import utils
+
+logger = logging.getLogger(__name__)
+
+
+def _try_import_strategy():
+    """Late import for Tesnorflow"""
+    from tensorflow.distribute.experimental import MultiWorkerMirroredStrategy
+    return MultiWorkerMirroredStrategy
+
+
+class TFRunner(object):
+    """Manages a TensorFlow model for training."""
+
+    def __init__(self, model_creator, data_creator, config=None,
+                 verbose=False):
+        """Initializes the runner.
+
+        Args:
+            model_creator (dict -> Model): see tf_trainer.py.
+            data_creator (dict -> tf.Dataset, tf.Dataset): see tf_trainer.py.
+            config (dict): see tf_trainer.py.
+            verbose (bool): Outputs training data if true.
+        """
+
+        self.model_creator = model_creator
+        self.data_creator = data_creator
+        self.config = {} if config is None else config
+        self.epoch = 0
+        self.verbose = verbose
+
+    def setup(self):
+        """Initializes the model."""
+        logger.debug("Creating dataset")
+        self.train_dataset, self.test_dataset = self.data_creator(self.config)
+
+        logger.debug("Creating model")
+        self.model = self.model_creator(self.config)
+
+    def setup_distributed(self, urls, world_rank, world_size):
+        """Sets up TensorFLow distributed environment and initializes the model.
+
+        Args:
+            urls (str): the URLs that each node uses to connect.
+            world_rank (int): the index of the runner.
+            world_size (int): the total number of runners.
+        """
+        assert len(urls) == world_size
+        tf_config = {
+            "cluster": {
+                "worker": urls
+            },
+            "task": {
+                "index": world_rank,
+                "type": "worker"
+            }
+        }
+        os.environ["TF_CONFIG"] = json.dumps(tf_config)
+
+        MultiWorkerMirroredStrategy = _try_import_strategy()
+        self.strategy = MultiWorkerMirroredStrategy()
+
+        self.train_dataset, self.test_dataset = self.data_creator(self.config)
+
+        logger.debug("Creating model with MultiWorkerMirroredStrategy")
+        with self.strategy.scope():
+            self.model = self.model_creator(self.config)
+
+        # For use in model.evaluate()
+        self.local_model = None
+
+    def step(self):
+        """Runs a training epoch and updates the model parameters."""
+        fit_default_config = {"verbose": self.verbose}
+        fit_default_config.update(self.config.get("fit_config", {}))
+
+        history = self.model.fit(self.train_dataset, **fit_default_config)
+        if history is None:
+            stats = {}
+        else:
+            stats = {"train_" + k: v[-1] for k, v in history.history.items()}
+
+        self.epoch += 1
+        return stats
+
+    def validate(self):
+        """Evaluates the model on the validation data set."""
+        stats = {}
+        evaluate_config = {"verbose": self.verbose}
+        evaluate_config.update(self.config.get("evaluate_config", {}))
+
+        results = self.model.evaluate(self.test_dataset, **evaluate_config)
+        if results is None:
+            # Using local Model since model.evaluate() returns None
+            # for MultiWorkerMirroredStrategy
+            logger.warning("Running a local model to get validation score.")
+            self.local_model = self.model_creator(self.config)
+            self.local_model.set_weights(self.model.get_weights())
+            results = self.local_model.evaluate(self.test_dataset,
+                                                **evaluate_config)
+
+        if isinstance(results, list):
+            stats = {
+                "validation_" + k: v
+                for k, v in zip(self.model.metrics_names, results)
+            }
+        else:
+            stats = {"loss": results}
+
+        return stats
+
+    def get_state(self):
+        """Returns the state of the runner."""
+        return {
+            "epoch": self.epoch,
+            "weights": self.model.get_weights(),
+            "optimizer_weights": self.model.optimizer.get_weights()
+        }
+
+    def set_state(self, state):
+        """Sets the state of the model."""
+
+        self.model = self.model_creator(self.config)
+        self.epoch = state["epoch"]
+        self.model.set_weights(state["weights"])
+        # This part is due to ray.get() changing scalar np.int64 object to int
+        state["optimizer_weights"][0] = np.array(
+            state["optimizer_weights"][0], dtype=np.int64)
+
+        if self.model.optimizer.weights == []:
+            self.model._make_train_function()
+        self.model.optimizer.set_weights(state["optimizer_weights"])
+
+    def shutdown(self):
+        """Attempts to shut down the worker."""
+        del self.model
+        del self.train_dataset
+        del self.test_dataset
+
+    def get_node_ip(self):
+        """Returns the IP address of the current node."""
+        return ray.services.get_node_ip_address()
+
+    def find_free_port(self):
+        """Finds a free port on the current node."""
+        return utils.find_free_port()

python/ray/experimental/sgd/tf/tf_trainer.py

@@ -0,0 +1,196 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import os
+import logging
+import pickle
+
+import ray
+
+from ray.tune import Trainable
+from ray.tune.resources import Resources
+from ray.experimental.sgd.tf.tf_runner import TFRunner
+
+logger = logging.getLogger(__name__)
+
+
+class TFTrainer(object):
+    def __init__(self,
+                 model_creator,
+                 data_creator,
+                 config=None,
+                 num_replicas=1,
+                 use_gpu=False,
+                 verbose=False):
+        """Sets up the TensorFlow trainer.
+
+        Args:
+            model_creator (dict -> Model): This function takes in the `config`
+                dict and returns a compiled TF model.
+            data_creator (dict -> tf.Dataset, tf.Dataset): Creates
+                the training and validation data sets using the config.
+                `config` dict is passed into the function.
+            config (dict): configuration passed to 'model_creator',
+                'data_creator'. Also contains `fit_config`, which is passed
+                into `model.fit(data, **fit_config)` and
+                `evaluate_config` which is passed into `model.evaluate`.
+            num_replicas (int): Sets number of workers used in distributed
+                training. Workers will be placed arbitrarily across the
+                cluster.
+            use_gpu (bool): Enables all workers to use GPU.
+            verbose (bool): Prints output of one model if true.
+        """
+        self.model_creator = model_creator
+        self.data_creator = data_creator
+        self.config = {} if config is None else config
+        self.use_gpu = use_gpu
+        self.num_replicas = num_replicas
+        self.verbose = verbose
+
+        # Generate actor class
+        Runner = ray.remote(num_cpus=1, num_gpus=int(use_gpu))(TFRunner)
+
+        if num_replicas == 1:
+            # Start workers
+            self.workers = [
+                Runner.remote(
+                    model_creator,
+                    data_creator,
+                    config=self.config,
+                    verbose=self.verbose)
+            ]
+            # Get setup tasks in order to throw errors on failure
+            ray.get(self.workers[0].setup.remote())
+        else:
+            # Start workers
+            self.workers = [
+                Runner.remote(
+                    model_creator,
+                    data_creator,
+                    config=self.config,
+                    verbose=self.verbose and i == 0)
+                for i in range(num_replicas)
+            ]
+
+            # Compute URL for initializing distributed setup
+            ips = ray.get(
+                [worker.get_node_ip.remote() for worker in self.workers])
+            ports = ray.get(
+                [worker.find_free_port.remote() for worker in self.workers])
+
+            urls = [
+                "{ip}:{port}".format(ip=ips[i], port=ports[i])
+                for i in range(len(self.workers))
+            ]
+
+            # Get setup tasks in order to throw errors on failure
+            ray.get([
+                worker.setup_distributed.remote(urls, i, len(self.workers))
+                for i, worker in enumerate(self.workers)
+            ])
+
+    def train(self):
+        """Runs a training epoch."""
+        worker_stats = ray.get([w.step.remote() for w in self.workers])
+        stats = worker_stats[0].copy()
+        return stats
+
+    def validate(self):
+        """Evaluates the model on the validation data set."""
+        logger.info("Starting validation step.")
+        stats = ray.get([w.validate.remote() for w in self.workers])
+        stats = stats[0].copy()
+        return stats
+
+    def get_model(self):
+        """Returns the learned model."""
+        state = ray.get(self.workers[0].get_state.remote())
+        return self._get_model_from_state(state)
+
+    def save(self, checkpoint):
+        """Saves the model at the provided checkpoint.
+
+        Args:
+            checkpoint (str): Path to target checkpoint file.
+
+        """
+
+        state = ray.get(self.workers[0].get_state.remote())
+
+        with open(checkpoint, "wb") as f:
+            pickle.dump(state, f)
+
+        return checkpoint
+
+    def restore(self, checkpoint):
+        """Restores the model from the provided checkpoint.
+
+        Args:
+            checkpoint (str): Path to target checkpoint file.
+
+        """
+        with open(checkpoint, "rb") as f:
+            state = pickle.load(f)
+
+        state_id = ray.put(state)
+        ray.get([worker.set_state.remote(state_id) for worker in self.workers])
+
+    def shutdown(self):
+        """Shuts down workers and releases resources."""
+        for worker in self.workers:
+            worker.shutdown.remote()
+            worker.__ray_terminate__.remote()
+
+    def _get_model_from_state(self, state):
+        """Creates model and load weights from state"""
+
+        model = self.model_creator(self.config)
+        model.set_weights(state["weights"])
+
+        # This part is due to ray.get() changing scalar np.int64 object to int
+        state["optimizer_weights"][0] = np.array(
+            state["optimizer_weights"][0], dtype=np.int64)
+
+        if model.optimizer.weights == []:
+            model._make_train_function()
+        model.optimizer.set_weights(state["optimizer_weights"])
+
+        return model
+
+
+class TFTrainable(Trainable):
+    @classmethod
+    def default_resource_request(cls, config):
+        return Resources(
+            cpu=0,
+            gpu=0,
+            extra_cpu=config["num_replicas"],
+            extra_gpu=int(config["use_gpu"]) * config["num_replicas"])
+
+    def _setup(self, config):
+        self._trainer = TFTrainer(
+            model_creator=config["model_creator"],
+            data_creator=config["data_creator"],
+            config=config.get("trainer_config", {}),
+            num_replicas=config["num_replicas"],
+            use_gpu=config["use_gpu"])
+
+    def _train(self):
+
+        train_stats = self._trainer.train()
+        validation_stats = self._trainer.validate()
+
+        train_stats.update(validation_stats)
+
+        return train_stats
+
+    def _save(self, checkpoint_dir):
+        return self._trainer.save(os.path.join(checkpoint_dir, "model"))
+
+    def _restore(self, checkpoint_path):
+        return self._trainer.restore(checkpoint_path)
+
+    def _stop(self):
+        self._trainer.shutdown()

python/ray/experimental/sgd/utils.py

@@ -6,92 +6,6 @@ from contextlib import closing
 import numpy as np
 import socket
 import time
-import torch
-import torch.nn as nn
-
-
-def train(train_iterator, model, criterion, optimizer):
-    """Runs 1 training epoch"""
-    batch_time = AverageMeter()
-    data_time = AverageMeter()
-    losses = AverageMeter()
-
-    timers = {k: TimerStat() for k in ["d2h", "fwd", "grad", "apply"]}
-
-    # switch to train mode
-    model.train()
-
-    end = time.time()
-
-    for i, (features, target) in enumerate(train_iterator):
-        # measure data loading time
-        data_time.update(time.time() - end)
-
-        # Create non_blocking tensors for distributed training
-        with timers["d2h"]:
-            if torch.cuda.is_available():
-                features = features.cuda(non_blocking=True)
-                target = target.cuda(non_blocking=True)
-
-        # compute output
-        with timers["fwd"]:
-            output = model(features)
-            loss = criterion(output, target)
-
-            # measure accuracy and record loss
-            losses.update(loss.item(), features.size(0))
-
-        with timers["grad"]:
-            # compute gradients in a backward pass
-            optimizer.zero_grad()
-            loss.backward()
-
-        with timers["apply"]:
-            # Call step of optimizer to update model params
-            optimizer.step()
-
-        # measure elapsed time
-        batch_time.update(time.time() - end)
-        end = time.time()
-
-    stats = {
-        "batch_time": batch_time.avg,
-        "batch_processed": losses.count,
-        "train_loss": losses.avg,
-        "data_time": data_time.avg,
-    }
-    stats.update({k: t.mean for k, t in timers.items()})
-    return stats
-
-
-def validate(val_loader, model, criterion):
-    batch_time = AverageMeter()
-    losses = AverageMeter()
-
-    # switch to evaluate mode
-    model.eval()
-
-    with torch.no_grad():
-        end = time.time()
-        for i, (features, target) in enumerate(val_loader):
-
-            if torch.cuda.is_available():
-                features = features.cuda(non_blocking=True)
-                target = target.cuda(non_blocking=True)
-
-            # compute output
-            output = model(features)
-            loss = criterion(output, target)
-
-            # measure accuracy and record loss
-            losses.update(loss.item(), features.size(0))
-
-            # measure elapsed time
-            batch_time.update(time.time() - end)
-            end = time.time()
-
-    stats = {"batch_time": batch_time.avg, "validation_loss": losses.avg}
-    return stats
 
 
 class TimerStat(object):
@@ -211,17 +125,3 @@ class AverageMeter(object):
         self.sum += val * n
         self.count += n
         self.avg = self.sum / self.count
-
-
-def sgd_mse_optimizer(model, config):
-    """Returns the mean squared error criterion and SGD optimizer.
-
-    Args:
-        model (torch.nn.Module): the model to optimize.
-        config (dict): configuration for the optimizer.
-            lr (float): the learning rate. defaults to 0.01.
-    """
-    learning_rate = config.get("lr", 0.01)
-    criterion = nn.MSELoss()
-    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
-    return criterion, optimizer