Ensure MXNet notebooks run in distributed mode.

Improved the notebooks so that they split the training data.

Author: Ignacio Quintero

sagemaker-python-sdk/mxnet_gluon_cifar10/cifar10.py

@@ -36,8 +36,17 @@ def train(current_host, hosts, num_cpus, num_gpus, channel_input_dirs, model_dir
  # load training and validation data
  # we use the gluon.data.vision.CIFAR10 class because of its built in pre-processing logic,
  # but point it at the location where SageMaker placed the data files, so it doesn't download them again.
+
+ part_index = 0
+ for i, host in enumerate(hosts):
+ if host == current_host:
+ part_index = i
+ break
+
+
  data_dir = channel_input_dirs['training']
- train_data = get_train_data(num_cpus, data_dir, batch_size, (3, 32, 32))
+ train_data = get_train_data(num_cpus, data_dir, batch_size, (3, 32, 32),
+ num_parts=len(hosts), part_index=part_index)
  test_data = get_test_data(num_cpus, data_dir, batch_size, (3, 32, 32))
 
  # Collect all parameters from net and its children, then initialize them.
@@ -104,23 +113,26 @@ def save(net, model_dir):
  os.rename(os.path.join(model_dir, best), os.path.join(model_dir, 'model.params'))
 
 
-def get_data(path, augment, num_cpus, batch_size, data_shape, resize=-1):
+def get_data(path, augment, num_cpus, batch_size, data_shape, resize=-1, num_parts=1, part_index=0):
  return mx.io.ImageRecordIter(
  path_imgrec=path,
  resize=resize,
  data_shape=data_shape,
  batch_size=batch_size,
  rand_crop=augment,
  rand_mirror=augment,
- preprocess_threads=num_cpus)
+ preprocess_threads=num_cpus,
+ num_parts=num_parts,
+ part_index=part_index)
 
 
 def get_test_data(num_cpus, data_dir, batch_size, data_shape, resize=-1):
- return get_data(os.path.join(data_dir, "test.rec"), False, num_cpus, batch_size, data_shape, resize)
+ return get_data(os.path.join(data_dir, "test.rec"), False, num_cpus, batch_size, data_shape, resize, 1, 0)
 
 
-def get_train_data(num_cpus, data_dir, batch_size, data_shape, resize=-1):
- return get_data(os.path.join(data_dir, "train.rec"), True, num_cpus, batch_size, data_shape, resize)
+def get_train_data(num_cpus, data_dir, batch_size, data_shape, resize=-1, num_parts=1, part_index=0):
+ return get_data(os.path.join(data_dir, "train.rec"), True, num_cpus, batch_size, data_shape, resize, num_parts,
+ part_index)
 
 
 def test(ctx, net, test_data):

sagemaker-python-sdk/mxnet_gluon_mnist/mnist.py

@@ -16,7 +16,7 @@
 # ------------------------------------------------------------ #
 
 
-def train(channel_input_dirs, hyperparameters, hosts, num_gpus, **kwargs):
+def train(current_host, channel_input_dirs, hyperparameters, hosts, num_gpus):
  # SageMaker passes num_cpus, num_gpus and other args we can use to tailor training to
  # the current container environment, but here we just use simple cpu context.
  ctx = mx.cpu()
@@ -53,6 +53,19 @@ def train(channel_input_dirs, hyperparameters, hosts, num_gpus, **kwargs):
  metric = mx.metric.Accuracy()
  loss = gluon.loss.SoftmaxCrossEntropyLoss()
 
+ # shard the training data in case we are doing
+ # distributed training.
+ if len(hosts) > 1:
+ train_data = [x for x in train_data]
+ shard_size = len(train_data) / len(hosts)
+ for i, host in enumerate(hosts):
+ if host == current_host:
+ start = shard_size * i
+ end = start + shard_size
+ break
+
+ train_data = train_data[start:end]
+
  for epoch in range(epochs):
  # reset data iterator and metric at begining of epoch.
  metric.reset()

sagemaker-python-sdk/mxnet_gluon_sentiment/sentiment.py

@@ -46,7 +46,14 @@ def train(current_host, hosts, num_cpus, num_gpus, channel_input_dirs, model_dir
  train_sentences = [[vocab.get(token, 1) for token in line if len(line)>0] for line in train_sentences]
  val_sentences = [[vocab.get(token, 1) for token in line if len(line)>0] for line in val_sentences]
 
- train_iterator = BucketSentenceIter(train_sentences, train_labels, batch_size)
+ shard_size = len(train_sentences) / len(hosts)
+ for i, host in enumerate(hosts):
+ if host == current_host:
+ start = shard_size * i
+ end = start + shard_size
+ break
+
+ train_iterator = BucketSentenceIter(train_sentences[start:end], train_labels[start:end], batch_size)
  val_iterator = BucketSentenceIter(val_sentences, val_labels, batch_size)
 
  # define the network

sagemaker-python-sdk/mxnet_mnist/mnist.py

@@ -35,11 +35,19 @@ def build_graph():
  return mx.sym.SoftmaxOutput(data=fc3, name='softmax')
 
 
-def train(channel_input_dirs, hyperparameters, hosts, num_cpus, num_gpus, **kwargs):
+def train(current_host, channel_input_dirs, hyperparameters, hosts, num_cpus, num_gpus):
  (train_labels, train_images) = load_data(os.path.join(channel_input_dirs['train']))
  (test_labels, test_images) = load_data(os.path.join(channel_input_dirs['test']))
+
+ shard_size = len(train_images) / len(hosts)
+ for i, host in enumerate(hosts):
+ if host == current_host:
+ start = shard_size * i
+ end = start + shard_size
+ break
+
  batch_size = 100
- train_iter = mx.io.NDArrayIter(train_images, train_labels, batch_size, shuffle=True)
+ train_iter = mx.io.NDArrayIter(train_images[start:end], train_labels[start:end], batch_size, shuffle=True)
  val_iter = mx.io.NDArrayIter(test_images, test_labels, batch_size)
  logging.getLogger().setLevel(logging.DEBUG)
  kvstore = 'local' if len(hosts) == 1 else 'dist_sync'