def main():

    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)
    kernel_size = 5
    strides = 2
    pool_size = 2
    hidden_dim = 512
    output_dim = 10
    activation = 'relu'

    def conv_nets():
        net = gluon.nn.HybridSequential()
        with net.name_scope():
            net.add(gluon.nn.Conv2D(channels=20, kernel_size=kernel_size, activation=activation))
            net.add(gluon.nn.MaxPool2D(pool_size=pool_size, strides=strides))
            net.add(gluon.nn.Conv2D(channels=50, kernel_size=kernel_size, activation=activation))
            net.add(gluon.nn.MaxPool2D(pool_size=pool_size, strides=strides))
            net.add(gluon.nn.Flatten())
            net.add(gluon.nn.Dense(hidden_dim, activation=activation))
            net.add(gluon.nn.Dense(output_dim))
        return net

    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()
    hvd.init()
    context = (mx.cpu(hvd.local_rank()) if args.no_cuda else mx.gpu(hvd.local_rank()))
    num_workers = hvd.size()
    (train_data, val_data) = get_mnist_iterator(hvd.rank())
    model = conv_nets()
    model.cast(args.dtype)
    model.hybridize()
    optimizer_params = {'momentum': args.momentum, 'learning_rate': (args.lr * hvd.size())}
    opt = mx.optimizer.create('sgd', **optimizer_params)
    initializer = mx.init.Xavier(rnd_type='gaussian', factor_type='in', magnitude=2)
    model.initialize(initializer, ctx=context)
    params = model.collect_params()
    if (params is not None):
        hvd.broadcast_parameters(params, root_rank=0)
    trainer = hvd.DistributedTrainer(params, opt)
    loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()
    metric = mx.metric.Accuracy()
    if (hvd.rank() == 0):
        global_tic = time.time()
    for epoch in range(args.epochs):
        tic = time.time()
        train_data.reset()
        metric.reset()
        for (nbatch, batch) in enumerate(train_data, start=1):
            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)
        (_, 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)
    if (hvd.rank() == 0):
        global_training_time = (time.time() - global_tic)
        print('Global elpased time on training:{}'.format(global_training_time))
        device = (context.device_type + str(num_workers))
        logging.info('Device info: %s', device)
def main():

    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)
    kernel_size = 5
    strides = 2
    pool_size = 2
    hidden_dim = 512
    output_dim = 10
    activation = 'relu'

    def conv_nets():
        net = gluon.nn.HybridSequential()
        with net.name_scope():
            net.add(gluon.nn.Conv2D(channels=20, kernel_size=kernel_size, activation=activation))
            net.add(gluon.nn.MaxPool2D(pool_size=pool_size, strides=strides))
            net.add(gluon.nn.Conv2D(channels=50, kernel_size=kernel_size, activation=activation))
            net.add(gluon.nn.MaxPool2D(pool_size=pool_size, strides=strides))
            net.add(gluon.nn.Flatten())
            net.add(gluon.nn.Dense(hidden_dim, activation=activation))
            net.add(gluon.nn.Dense(output_dim))
        return net

    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()
    hvd.init()
    context = (mx.cpu(hvd.local_rank()) if args.no_cuda else mx.gpu(hvd.local_rank()))
    num_workers = hvd.size()
    (train_data, val_data) = get_mnist_iterator(hvd.rank())
    model = conv_nets()
    model.cast(args.dtype)
    model.hybridize()
    optimizer_params = {'momentum': args.momentum, 'learning_rate': (args.lr * hvd.size())}
    opt = mx.optimizer.create('sgd', **optimizer_params)
    initializer = mx.init.Xavier(rnd_type='gaussian', factor_type='in', magnitude=2)
    model.initialize(initializer, ctx=context)
    params = model.collect_params()
    if (params is not None):
        hvd.broadcast_parameters(params, root_rank=0)
    trainer = hvd.DistributedTrainer(params, opt)
    loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()
    metric = mx.metric.Accuracy()
    if (hvd.rank() == 0):
        global_tic = time.time()
    for epoch in range(args.epochs):
        tic = time.time()
        train_data.reset()
        metric.reset()
        for (nbatch, batch) in enumerate(train_data, start=1):
            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)
        (_, 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)
    if (hvd.rank() == 0):
        global_training_time = (time.time() - global_tic)
        print('Global elpased time on training:{}'.format(global_training_time))
        device = (context.device_type + str(num_workers))
        logging.info('Device info: %s', device)