Support Mxnet Horovod (aws#373)

* support mxnet horovod

* mxnet hvd example

Author: NihalHarish

examples/mxnet/scripts/mnist_mxnet_hvd.py

@@ -0,0 +1,201 @@
+import argparse
+import logging
+import os
+import zipfile
+import time
+
+import mxnet as mx
+import horovod.mxnet as hvd
+from mxnet import autograd, gluon, nd
+from mxnet.test_utils import download
+
+from tornasole import SaveConfig, modes
+from tornasole.mxnet import TornasoleHook
+
+# Training settings
+parser = argparse.ArgumentParser(description="MXNet MNIST Example")
+
+parser.add_argument("--batch-size", type=int, default=64, help="training batch size (default: 64)")
+parser.add_argument(
+    "--dtype", type=str, default="float32", help="training data type (default: float32)"
+)
+parser.add_argument("--epochs", type=int, default=5, help="number of training epochs (default: 5)")
+parser.add_argument("--lr", type=float, default=0.01, help="learning rate (default: 0.01)")
+parser.add_argument("--momentum", type=float, default=0.9, help="SGD momentum (default: 0.9)")
+parser.add_argument(
+    "--no-cuda", action="store_true", default=False, help="disable training on GPU (default: False)"
+)
+parser.add_argument(
+    "--output-uri",
+    type=str,
+    default="/opt/ml/output/tensors/tornasole",
+    help="S3 URI of the bucket where tensor data will be stored.",
+)
+args = parser.parse_args()
+
+if not args.no_cuda:
+    # Disable CUDA if there are no GPUs.
+    if not mx.test_utils.list_gpus():
+        args.no_cuda = True
+
+logging.basicConfig(level=logging.INFO)
+logging.info(args)
+
+
+# Function to get mnist iterator given a rank
+def get_mnist_iterator(rank):
+    data_dir = "data-%d" % rank
+    if not os.path.isdir(data_dir):
+        os.makedirs(data_dir)
+    zip_file_path = download("http://data.mxnet.io/mxnet/data/mnist.zip", dirname=data_dir)
+    with zipfile.ZipFile(zip_file_path) as zf:
+        zf.extractall(data_dir)
+
+    input_shape = (1, 28, 28)
+    batch_size = args.batch_size
+
+    train_iter = mx.io.MNISTIter(
+        image="%s/train-images-idx3-ubyte" % data_dir,
+        label="%s/train-labels-idx1-ubyte" % data_dir,
+        input_shape=input_shape,
+        batch_size=batch_size,
+        shuffle=True,
+        flat=False,
+        num_parts=hvd.size(),
+        part_index=hvd.rank(),
+    )
+
+    val_iter = mx.io.MNISTIter(
+        image="%s/t10k-images-idx3-ubyte" % data_dir,
+        label="%s/t10k-labels-idx1-ubyte" % data_dir,
+        input_shape=input_shape,
+        batch_size=batch_size,
+        flat=False,
+    )
+
+    return train_iter, val_iter
+
+
+# Function to define neural network
+def conv_nets():
+    net = gluon.nn.HybridSequential()
+    with net.name_scope():
+        net.add(gluon.nn.Conv2D(channels=20, kernel_size=5, activation="relu"))
+        net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))
+        net.add(gluon.nn.Conv2D(channels=50, kernel_size=5, activation="relu"))
+        net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))
+        net.add(gluon.nn.Flatten())
+        net.add(gluon.nn.Dense(512, activation="relu"))
+        net.add(gluon.nn.Dense(10))
+    return net
+
+
+# Function to evaluate accuracy for a model
+def evaluate(model, data_iter, context):
+    data_iter.reset()
+    metric = mx.metric.Accuracy()
+    for _, batch in enumerate(data_iter):
+        data = batch.data[0].as_in_context(context)
+        label = batch.label[0].as_in_context(context)
+        output = model(data.astype(args.dtype, copy=False))
+        metric.update([label], [output])
+
+    return metric.get()
+
+
+# Initialize Horovod
+hvd.init()
+
+# Horovod: pin context to local rank
+context = mx.cpu(hvd.local_rank()) if args.no_cuda else mx.gpu(hvd.local_rank())
+num_workers = hvd.size()
+
+# Load training and validation data
+train_data, val_data = get_mnist_iterator(hvd.rank())
+
+# Build model
+model = conv_nets()
+model.cast(args.dtype)
+model.hybridize()
+
+# Create optimizer
+optimizer_params = {"momentum": args.momentum, "learning_rate": args.lr * hvd.size()}
+opt = mx.optimizer.create("sgd", **optimizer_params)
+# opt = ts.TornasoleOptimizer(opt)
+
+# Initialize parameters
+initializer = mx.init.Xavier(rnd_type="gaussian", factor_type="in", magnitude=2)
+model.initialize(initializer, ctx=context)
+
+# Horovod: fetch and broadcast parameters
+params = model.collect_params()
+if params is not None:
+    hvd.broadcast_parameters(params, root_rank=0)
+
+# Horovod: create DistributedTrainer, a subclass of gluon.Trainer
+trainer = hvd.DistributedTrainer(params, opt)
+
+# Create loss function and train metric
+loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()
+metric = mx.metric.Accuracy()
+
+
+def create_tornasole_hook():
+    # With the following SaveConfig, we will save tensors for steps 1, 2 and 3
+    # (indexing starts with 0).
+    save_config = SaveConfig(save_interval=1)
+
+    # Create a hook that logs weights, biases and gradients while training the model.
+    ts_hook = TornasoleHook(
+        out_dir=args.output_uri,
+        save_config=save_config,
+        include_collections=["weights", "gradients", "biases"],
+    )
+    return ts_hook
+
+
+# Train model
+for epoch in range(args.epochs):
+    tic = time.time()
+    train_data.reset()
+    metric.reset()
+
+    # Create Tornasole Hook
+    hook = create_tornasole_hook()
+    hook.register_hook(model)
+
+    for nbatch, batch in enumerate(train_data, start=1):
+        hook.set_mode(modes.TRAIN)
+        data = batch.data[0].as_in_context(context)
+        label = batch.label[0].as_in_context(context)
+        with autograd.record():
+            output = model(data.astype(args.dtype, copy=False))
+            loss = loss_fn(output, label)
+        loss.backward()
+        trainer.step(args.batch_size)
+        metric.update([label], [output])
+
+        if nbatch % 100 == 0:
+            name, acc = metric.get()
+            logging.info("[Epoch %d Batch %d] Training: %s=%f" % (epoch, nbatch, name, acc))
+
+    if hvd.rank() == 0:
+        elapsed = time.time() - tic
+        speed = nbatch * args.batch_size * hvd.size() / elapsed
+        logging.info("Epoch[%d]\tSpeed=%.2f samples/s\tTime cost=%f", epoch, speed, elapsed)
+
+    # Evaluate model accuracy
+    hook.set_mode(modes.EVAL)
+    _, train_acc = metric.get()
+    name, val_acc = evaluate(model, val_data, context)
+    if hvd.rank() == 0:
+        logging.info(
+            "Epoch[%d]\tTrain: %s=%f\tValidation: %s=%f", epoch, name, train_acc, name, val_acc
+        )
+
+    if hvd.rank() == 0 and epoch == args.epochs - 1:
+        assert val_acc > 0.96, (
+            "Achieved accuracy (%f) is lower than expected\
+                                (0.96)"
+            % val_acc
+        )

tornasole/mxnet/hook.py

@@ -59,9 +59,23 @@ def __init__(
         set_hook(self)
 
     def get_worker_name(self):
+        try:
+            import horovod.mxnet as hvd
+
+            if hvd.size():
+                return f"worker_{hvd.rank()}"
+        except (ModuleNotFoundError, ValueError, ImportError):
+            pass
         return CONFIG_DEFAULT_WORKER_NAME
 
     def get_num_workers(self):
+        try:
+            import horovod.mxnet as hvd
+
+            if hvd.size():
+                return hvd.size()
+        except (ModuleNotFoundError, ValueError, ImportError):
+            pass
         return 1
 
     @classmethod