We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
mxnet+hvd with different random seed won't work with Gluon API
**Environment:**
1. Framework: (MXNet)
2. Framework version: 1.4
3. Horovod version: v0.16.0 
4. MPI version: 
5. CUDA version:
6. NCCL version:
7. Python version: 2.7/3.6
8. OS and version: 

**Checklist:**
1. Did you search issues to find if somebody asked this question before?
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.md)?
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.md)?

**Bug report:**
Please describe errorneous behavior you're observing and steps to reproduce it.

If the shape of a parameter has to be inferred after a batch of data is seen, the parameter will not be broadcast. Therefore, the correctness of the program depends on whether all workers are initialized with the same random seed. So the following program will not work:

```
random.seed(hvd.local_rank())
data = ..
model = mx.gluon.nn.Dense(10)
model.initialize()

// no parameters are broadcast, because initialization is deferred.
hvd.broadcast_parameters(model.collect_parameters())

for batch in data:
    // params are initialized after shape is known
    pred = model(batch.data[0])
    ...
...
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/mxnet_mnist.py`
Content:
```
1 import argparse
2 import logging
3 import os
4 import zipfile
5 import time
6 
7 import mxnet as mx
8 import horovod.mxnet as hvd
9 from mxnet import autograd, gluon, nd
10 from mxnet.test_utils import download
11 
12 # Training settings
13 parser = argparse.ArgumentParser(description='MXNet MNIST Example')
14 
15 parser.add_argument('--batch-size', type=int, default=64,
16                     help='training batch size (default: 64)')
17 parser.add_argument('--dtype', type=str, default='float32',
18                     help='training data type (default: float32)')
19 parser.add_argument('--epochs', type=int, default=5,
20                     help='number of training epochs (default: 5)')
21 parser.add_argument('--lr', type=float, default=0.01,
22                     help='learning rate (default: 0.01)')
23 parser.add_argument('--momentum', type=float, default=0.9,
24                     help='SGD momentum (default: 0.9)')
25 parser.add_argument('--no-cuda', action='store_true', default=False,
26                     help='disable training on GPU (default: False)')
27 args = parser.parse_args()
28 
29 if not args.no_cuda:
30     # Disable CUDA if there are no GPUs.
31     if not mx.test_utils.list_gpus():
32         args.no_cuda = True
33 
34 logging.basicConfig(level=logging.INFO)
35 logging.info(args)
36 
37 
38 # Function to get mnist iterator given a rank
39 def get_mnist_iterator(rank):
40     data_dir = "data-%d" % rank
41     if not os.path.isdir(data_dir):
42         os.makedirs(data_dir)
43     zip_file_path = download('http://data.mxnet.io/mxnet/data/mnist.zip',
44                              dirname=data_dir)
45     with zipfile.ZipFile(zip_file_path) as zf:
46         zf.extractall(data_dir)
47 
48     input_shape = (1, 28, 28)
49     batch_size = args.batch_size
50 
51     train_iter = mx.io.MNISTIter(
52         image="%s/train-images-idx3-ubyte" % data_dir,
53         label="%s/train-labels-idx1-ubyte" % data_dir,
54         input_shape=input_shape,
55         batch_size=batch_size,
56         shuffle=True,
57         flat=False,
58         num_parts=hvd.size(),
59         part_index=hvd.rank()
60     )
61 
62     val_iter = mx.io.MNISTIter(
63         image="%s/t10k-images-idx3-ubyte" % data_dir,
64         label="%s/t10k-labels-idx1-ubyte" % data_dir,
65         input_shape=input_shape,
66         batch_size=batch_size,
67         flat=False,
68     )
69 
70     return train_iter, val_iter
71 
72 
73 # Function to define neural network
74 def conv_nets():
75     net = gluon.nn.HybridSequential()
76     with net.name_scope():
77         net.add(gluon.nn.Conv2D(channels=20, kernel_size=5, activation='relu'))
78         net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))
79         net.add(gluon.nn.Conv2D(channels=50, kernel_size=5, activation='relu'))
80         net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))
81         net.add(gluon.nn.Flatten())
82         net.add(gluon.nn.Dense(512, activation="relu"))
83         net.add(gluon.nn.Dense(10))
84     return net
85 
86 
87 # Function to evaluate accuracy for a model
88 def evaluate(model, data_iter, context):
89     data_iter.reset()
90     metric = mx.metric.Accuracy()
91     for _, batch in enumerate(data_iter):
92         data = batch.data[0].as_in_context(context)
93         label = batch.label[0].as_in_context(context)
94         output = model(data.astype(args.dtype, copy=False))
95         metric.update([label], [output])
96 
97     return metric.get()
98 
99 
100 # Initialize Horovod
101 hvd.init()
102 
103 # Horovod: pin context to local rank
104 context = mx.cpu(hvd.local_rank()) if args.no_cuda else mx.gpu(hvd.local_rank())
105 num_workers = hvd.size()
106 
107 # Load training and validation data
108 train_data, val_data = get_mnist_iterator(hvd.rank())
109 
110 # Build model
111 model = conv_nets()
112 model.cast(args.dtype)
113 model.hybridize()
114 
115 # Define hyper parameters
116 optimizer_params = {'momentum': args.momentum,
117                     'learning_rate': args.lr * hvd.size(),
118                     'rescale_grad': 1.0 / args.batch_size}
119 
120 # Add Horovod Distributed Optimizer
121 opt = mx.optimizer.create('sgd', **optimizer_params)
122 opt = hvd.DistributedOptimizer(opt)
123 
124 # Initialize parameters
125 initializer = mx.init.Xavier(rnd_type='gaussian', factor_type="in",
126                              magnitude=2)
127 model.initialize(initializer, ctx=context)
128 
129 # Fetch and broadcast parameters
130 params = model.collect_params()
131 if params is not None:
132     hvd.broadcast_parameters(params, root_rank=0)
133 
134 # Create trainer, loss function and train metric
135 trainer = gluon.Trainer(params, opt, kvstore=None)
136 loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()
137 metric = mx.metric.Accuracy()
138 
139 # Train model
140 for epoch in range(args.epochs):
141     tic = time.time()
142     train_data.reset()
143     metric.reset()
144     for nbatch, batch in enumerate(train_data, start=1):
145         data = batch.data[0].as_in_context(context)
146         label = batch.label[0].as_in_context(context)
147         with autograd.record():
148             output = model(data.astype(args.dtype, copy=False))
149             loss = loss_fn(output, label)
150         loss.backward()
151         trainer.step(args.batch_size)
152         metric.update([label], [output])
153 
154         if nbatch % 100 == 0:
155             name, acc = metric.get()
156             logging.info('[Epoch %d Batch %d] Training: %s=%f' %
157                          (epoch, nbatch, name, acc))
158 
159     if hvd.rank() == 0:
160         elapsed = time.time() - tic
161         speed = nbatch * args.batch_size * hvd.size() / elapsed
162         logging.info('Epoch[%d]\tSpeed=%.2f samples/s\tTime cost=%f',
163                      epoch, speed, elapsed)
164 
165     # Evaluate model accuracy
166     _, train_acc = metric.get()
167     name, val_acc = evaluate(model, val_data, context)
168     if hvd.rank() == 0:
169         logging.info('Epoch[%d]\tTrain: %s=%f\tValidation: %s=%f', epoch, name,
170                      train_acc, name, val_acc)
171 
172     if hvd.rank() == 0 and epoch == args.epochs - 1:
173         assert val_acc > 0.96, "Achieved accuracy (%f) is lower than expected\
174                                 (0.96)" % val_acc
175
```

Path: `horovod/mxnet/__init__.py`
Content:
```
1 # Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 
16 from __future__ import absolute_import
17 from __future__ import division
18 from __future__ import print_function
19 
20 from horovod.common import check_extension
21 
22 check_extension('horovod.mxnet', 'HOROVOD_WITH_MXNET',
23                 __file__, 'mpi_lib')
24 
25 from horovod.mxnet.mpi_ops import allgather
26 from horovod.mxnet.mpi_ops import allreduce, allreduce_
27 from horovod.mxnet.mpi_ops import broadcast, broadcast_
28 from horovod.mxnet.mpi_ops import init, shutdown
29 from horovod.mxnet.mpi_ops import size, local_size, rank, local_rank
30 from horovod.mxnet.mpi_ops import mpi_threads_supported
31 
32 import mxnet as mx
33 
34 
35 # This is where Horovod's DistributedOptimizer wrapper for MXNet goes
36 class DistributedOptimizer(mx.optimizer.Optimizer):
37     def __init__(self, optimizer):
38         self._optimizer = optimizer
39 
40     def __getattr__(self, item):
41         return getattr(self._optimizer, item)
42 
43     def create_state_multi_precision(self, index, weight):
44         return self._optimizer.create_state_multi_precision(index, weight)
45 
46     def _do_allreduce(self, index, grad):
47         if isinstance(index, (tuple, list)):
48             for i in range(len(index)):
49                 allreduce_(grad[i], average=True, name=str(index[i]))
50         else:
51             allreduce_(grad, average=True, name=str(index))
52 
53     def update(self, index, weight, grad, state):
54         self._do_allreduce(index, grad)
55         self._optimizer.update(index, weight, grad, state)
56 
57     def update_multi_precision(self, index, weight, grad, state):
58         self._do_allreduce(index, grad)
59         self._optimizer.update_multi_precision(index, weight, grad, state)
60 
61     def set_learning_rate(self, lr):
62         self._optimizer.set_learning_rate(lr)
63 
64     def set_lr_mult(self, args_lr_mult):
65         self._optimizer.set_lr_mult(args_lr_mult)
66 
67     def set_wd_mult(self, args_wd_mult):
68         self._optimizer.set_wd_mult(args_wd_mult)
69 
70 
71 def broadcast_parameters(params, root_rank=0):
72     """
73     Broadcasts the parameters from root rank to all other processes.
74     Typical usage is to broadcast the `Module.get_params()` or the
75     `Block.collect_params()`.
76 
77     Arguments:
78         params: One of the following:
79             - dict of parameters to broadcast
80             - ParameterDict to broadcast
81         root_rank: The rank of the process from which parameters will be
82                    broadcasted to all other processes.
83     """
84     tensors = []
85     if isinstance(params, dict):
86         tensors = [p for _, p in sorted(params.items())]
87     elif isinstance(params, mx.gluon.parameter.ParameterDict):
88         for _, p in sorted(params.items()):
89             try:
90                 tensors.append(p.data())
91             except mx.gluon.parameter.DeferredInitializationError:
92                 # skip broadcasting deferred init param
93                 pass
94     else:
95         raise ValueError('invalid params of type: %s' % type(params))
96 
97     # Run broadcasts.
98     for i, tensor in enumerate(tensors):
99         broadcast_(tensor, root_rank, str(i))
100 
101     # Make sure tensors pushed to MXNet engine get processed such that all
102     # workers are synced before starting training.
103     for tensor in tensors:
104         tensor.wait_to_read()
105
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/mxnet_mnist.py b/examples/mxnet_mnist.py
--- a/examples/mxnet_mnist.py
+++ b/examples/mxnet_mnist.py
@@ -112,13 +112,12 @@
 model.cast(args.dtype)
 model.hybridize()
 
-# Define hyper parameters
+# Create optimizer
 optimizer_params = {'momentum': args.momentum,
                     'learning_rate': args.lr * hvd.size(),
                     'rescale_grad': 1.0 / args.batch_size}
-
-# Add Horovod Distributed Optimizer
 opt = mx.optimizer.create('sgd', **optimizer_params)
+# Horovod: wrap optimizer with DistributedOptimizer
 opt = hvd.DistributedOptimizer(opt)
 
 # Initialize parameters
diff --git a/horovod/mxnet/__init__.py b/horovod/mxnet/__init__.py
--- a/horovod/mxnet/__init__.py
+++ b/horovod/mxnet/__init__.py
@@ -30,6 +30,7 @@
 from horovod.mxnet.mpi_ops import mpi_threads_supported
 
 import mxnet as mx
+import types
 
 
 # This is where Horovod's DistributedOptimizer wrapper for MXNet goes
@@ -68,6 +69,16 @@
         self._optimizer.set_wd_mult(args_wd_mult)
 
 
+# Wrapper to inject Horovod broadcast after parameter initialization
+def _append_broadcast_init(param, root_rank):
+    init_impl = getattr(param, '_init_impl')
+    def wrapped_init_impl(self, *args, **kwargs):
+        init_impl(*args, **kwargs)
+        broadcast_(self.data(), root_rank=root_rank)
+        self.data().wait_to_read()
+    return wrapped_init_impl
+
+
 def broadcast_parameters(params, root_rank=0):
     """
     Broadcasts the parameters from root rank to all other processes.
@@ -89,8 +100,10 @@
             try:
                 tensors.append(p.data())
             except mx.gluon.parameter.DeferredInitializationError:
-                # skip broadcasting deferred init param
-                pass
+                # Inject wrapper method with post-initialization broadcast to
+                # handle parameters with deferred initialization
+                new_init = _append_broadcast_init(p, root_rank)
+                p._init_impl = types.MethodType(new_init, p)
     else:
         raise ValueError('invalid params of type: %s' % type(params))

{"golden_diff": "diff --git a/examples/mxnet_mnist.py b/examples/mxnet_mnist.py\n--- a/examples/mxnet_mnist.py\n+++ b/examples/mxnet_mnist.py\n@@ -112,13 +112,12 @@\n model.cast(args.dtype)\n model.hybridize()\n \n-# Define hyper parameters\n+# Create optimizer\n optimizer_params = {'momentum': args.momentum,\n                     'learning_rate': args.lr * hvd.size(),\n                     'rescale_grad': 1.0 / args.batch_size}\n-\n-# Add Horovod Distributed Optimizer\n opt = mx.optimizer.create('sgd', **optimizer_params)\n+# Horovod: wrap optimizer with DistributedOptimizer\n opt = hvd.DistributedOptimizer(opt)\n \n # Initialize parameters\ndiff --git a/horovod/mxnet/__init__.py b/horovod/mxnet/__init__.py\n--- a/horovod/mxnet/__init__.py\n+++ b/horovod/mxnet/__init__.py\n@@ -30,6 +30,7 @@\n from horovod.mxnet.mpi_ops import mpi_threads_supported\n \n import mxnet as mx\n+import types\n \n \n # This is where Horovod's DistributedOptimizer wrapper for MXNet goes\n@@ -68,6 +69,16 @@\n         self._optimizer.set_wd_mult(args_wd_mult)\n \n \n+# Wrapper to inject Horovod broadcast after parameter initialization\n+def _append_broadcast_init(param, root_rank):\n+    init_impl = getattr(param, '_init_impl')\n+    def wrapped_init_impl(self, *args, **kwargs):\n+        init_impl(*args, **kwargs)\n+        broadcast_(self.data(), root_rank=root_rank)\n+        self.data().wait_to_read()\n+    return wrapped_init_impl\n+\n+\n def broadcast_parameters(params, root_rank=0):\n     \"\"\"\n     Broadcasts the parameters from root rank to all other processes.\n@@ -89,8 +100,10 @@\n             try:\n                 tensors.append(p.data())\n             except mx.gluon.parameter.DeferredInitializationError:\n-                # skip broadcasting deferred init param\n-                pass\n+                # Inject wrapper method with post-initialization broadcast to\n+                # handle parameters with deferred initialization\n+                new_init = _append_broadcast_init(p, root_rank)\n+                p._init_impl = types.MethodType(new_init, p)\n     else:\n         raise ValueError('invalid params of type: %s' % type(params))\n", "issue": "mxnet+hvd with different random seed won't work with Gluon API\n**Environment:**\r\n1. Framework: (MXNet)\r\n2. Framework version: 1.4\r\n3. Horovod version: v0.16.0 \r\n4. MPI version: \r\n5. CUDA version:\r\n6. NCCL version:\r\n7. Python version: 2.7/3.6\r\n8. OS and version: \r\n\r\n**Checklist:**\r\n1. Did you search issues to find if somebody asked this question before?\r\n2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.md)?\r\n3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.md)?\r\n\r\n**Bug report:**\r\nPlease describe errorneous behavior you're observing and steps to reproduce it.\r\n\r\nIf the shape of a parameter has to be inferred after a batch of data is seen, the parameter will not be broadcast. Therefore, the correctness of the program depends on whether all workers are initialized with the same random seed. So the following program will not work:\r\n\r\n```\r\nrandom.seed(hvd.local_rank())\r\ndata = ..\r\nmodel = mx.gluon.nn.Dense(10)\r\nmodel.initialize()\r\n\r\n// no parameters are broadcast, because initialization is deferred.\r\nhvd.broadcast_parameters(model.collect_parameters())\r\n\r\nfor batch in data:\r\n    // params are initialized after shape is known\r\n    pred = model(batch.data[0])\r\n    ...\r\n...\r\n```\r\n\n", "before_files": [{"content": "import argparse\nimport logging\nimport os\nimport zipfile\nimport time\n\nimport mxnet as mx\nimport horovod.mxnet as hvd\nfrom mxnet import autograd, gluon, nd\nfrom mxnet.test_utils import download\n\n# Training settings\nparser = argparse.ArgumentParser(description='MXNet MNIST Example')\n\nparser.add_argument('--batch-size', type=int, default=64,\n                    help='training batch size (default: 64)')\nparser.add_argument('--dtype', type=str, default='float32',\n                    help='training data type (default: float32)')\nparser.add_argument('--epochs', type=int, default=5,\n                    help='number of training epochs (default: 5)')\nparser.add_argument('--lr', type=float, default=0.01,\n                    help='learning rate (default: 0.01)')\nparser.add_argument('--momentum', type=float, default=0.9,\n                    help='SGD momentum (default: 0.9)')\nparser.add_argument('--no-cuda', action='store_true', default=False,\n                    help='disable training on GPU (default: False)')\nargs = parser.parse_args()\n\nif not args.no_cuda:\n    # Disable CUDA if there are no GPUs.\n    if not mx.test_utils.list_gpus():\n        args.no_cuda = True\n\nlogging.basicConfig(level=logging.INFO)\nlogging.info(args)\n\n\n# Function to get mnist iterator given a rank\ndef get_mnist_iterator(rank):\n    data_dir = \"data-%d\" % rank\n    if not os.path.isdir(data_dir):\n        os.makedirs(data_dir)\n    zip_file_path = download('http://data.mxnet.io/mxnet/data/mnist.zip',\n                             dirname=data_dir)\n    with zipfile.ZipFile(zip_file_path) as zf:\n        zf.extractall(data_dir)\n\n    input_shape = (1, 28, 28)\n    batch_size = args.batch_size\n\n    train_iter = mx.io.MNISTIter(\n        image=\"%s/train-images-idx3-ubyte\" % data_dir,\n        label=\"%s/train-labels-idx1-ubyte\" % data_dir,\n        input_shape=input_shape,\n        batch_size=batch_size,\n        shuffle=True,\n        flat=False,\n        num_parts=hvd.size(),\n        part_index=hvd.rank()\n    )\n\n    val_iter = mx.io.MNISTIter(\n        image=\"%s/t10k-images-idx3-ubyte\" % data_dir,\n        label=\"%s/t10k-labels-idx1-ubyte\" % data_dir,\n        input_shape=input_shape,\n        batch_size=batch_size,\n        flat=False,\n    )\n\n    return train_iter, val_iter\n\n\n# Function to define neural network\ndef conv_nets():\n    net = gluon.nn.HybridSequential()\n    with net.name_scope():\n        net.add(gluon.nn.Conv2D(channels=20, kernel_size=5, activation='relu'))\n        net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))\n        net.add(gluon.nn.Conv2D(channels=50, kernel_size=5, activation='relu'))\n        net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))\n        net.add(gluon.nn.Flatten())\n        net.add(gluon.nn.Dense(512, activation=\"relu\"))\n        net.add(gluon.nn.Dense(10))\n    return net\n\n\n# Function to evaluate accuracy for a model\ndef evaluate(model, data_iter, context):\n    data_iter.reset()\n    metric = mx.metric.Accuracy()\n    for _, batch in enumerate(data_iter):\n        data = batch.data[0].as_in_context(context)\n        label = batch.label[0].as_in_context(context)\n        output = model(data.astype(args.dtype, copy=False))\n        metric.update([label], [output])\n\n    return metric.get()\n\n\n# Initialize Horovod\nhvd.init()\n\n# Horovod: pin context to local rank\ncontext = mx.cpu(hvd.local_rank()) if args.no_cuda else mx.gpu(hvd.local_rank())\nnum_workers = hvd.size()\n\n# Load training and validation data\ntrain_data, val_data = get_mnist_iterator(hvd.rank())\n\n# Build model\nmodel = conv_nets()\nmodel.cast(args.dtype)\nmodel.hybridize()\n\n# Define hyper parameters\noptimizer_params = {'momentum': args.momentum,\n                    'learning_rate': args.lr * hvd.size(),\n                    'rescale_grad': 1.0 / args.batch_size}\n\n# Add Horovod Distributed Optimizer\nopt = mx.optimizer.create('sgd', **optimizer_params)\nopt = hvd.DistributedOptimizer(opt)\n\n# Initialize parameters\ninitializer = mx.init.Xavier(rnd_type='gaussian', factor_type=\"in\",\n                             magnitude=2)\nmodel.initialize(initializer, ctx=context)\n\n# Fetch and broadcast parameters\nparams = model.collect_params()\nif params is not None:\n    hvd.broadcast_parameters(params, root_rank=0)\n\n# Create trainer, loss function and train metric\ntrainer = gluon.Trainer(params, opt, kvstore=None)\nloss_fn = gluon.loss.SoftmaxCrossEntropyLoss()\nmetric = mx.metric.Accuracy()\n\n# Train model\nfor epoch in range(args.epochs):\n    tic = time.time()\n    train_data.reset()\n    metric.reset()\n    for nbatch, batch in enumerate(train_data, start=1):\n        data = batch.data[0].as_in_context(context)\n        label = batch.label[0].as_in_context(context)\n        with autograd.record():\n            output = model(data.astype(args.dtype, copy=False))\n            loss = loss_fn(output, label)\n        loss.backward()\n        trainer.step(args.batch_size)\n        metric.update([label], [output])\n\n        if nbatch % 100 == 0:\n            name, acc = metric.get()\n            logging.info('[Epoch %d Batch %d] Training: %s=%f' %\n                         (epoch, nbatch, name, acc))\n\n    if hvd.rank() == 0:\n        elapsed = time.time() - tic\n        speed = nbatch * args.batch_size * hvd.size() / elapsed\n        logging.info('Epoch[%d]\\tSpeed=%.2f samples/s\\tTime cost=%f',\n                     epoch, speed, elapsed)\n\n    # Evaluate model accuracy\n    _, train_acc = metric.get()\n    name, val_acc = evaluate(model, val_data, context)\n    if hvd.rank() == 0:\n        logging.info('Epoch[%d]\\tTrain: %s=%f\\tValidation: %s=%f', epoch, name,\n                     train_acc, name, val_acc)\n\n    if hvd.rank() == 0 and epoch == args.epochs - 1:\n        assert val_acc > 0.96, \"Achieved accuracy (%f) is lower than expected\\\n                                (0.96)\" % val_acc\n", "path": "examples/mxnet_mnist.py"}, {"content": "# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom horovod.common import check_extension\n\ncheck_extension('horovod.mxnet', 'HOROVOD_WITH_MXNET',\n                __file__, 'mpi_lib')\n\nfrom horovod.mxnet.mpi_ops import allgather\nfrom horovod.mxnet.mpi_ops import allreduce, allreduce_\nfrom horovod.mxnet.mpi_ops import broadcast, broadcast_\nfrom horovod.mxnet.mpi_ops import init, shutdown\nfrom horovod.mxnet.mpi_ops import size, local_size, rank, local_rank\nfrom horovod.mxnet.mpi_ops import mpi_threads_supported\n\nimport mxnet as mx\n\n\n# This is where Horovod's DistributedOptimizer wrapper for MXNet goes\nclass DistributedOptimizer(mx.optimizer.Optimizer):\n    def __init__(self, optimizer):\n        self._optimizer = optimizer\n\n    def __getattr__(self, item):\n        return getattr(self._optimizer, item)\n\n    def create_state_multi_precision(self, index, weight):\n        return self._optimizer.create_state_multi_precision(index, weight)\n\n    def _do_allreduce(self, index, grad):\n        if isinstance(index, (tuple, list)):\n            for i in range(len(index)):\n                allreduce_(grad[i], average=True, name=str(index[i]))\n        else:\n            allreduce_(grad, average=True, name=str(index))\n\n    def update(self, index, weight, grad, state):\n        self._do_allreduce(index, grad)\n        self._optimizer.update(index, weight, grad, state)\n\n    def update_multi_precision(self, index, weight, grad, state):\n        self._do_allreduce(index, grad)\n        self._optimizer.update_multi_precision(index, weight, grad, state)\n\n    def set_learning_rate(self, lr):\n        self._optimizer.set_learning_rate(lr)\n\n    def set_lr_mult(self, args_lr_mult):\n        self._optimizer.set_lr_mult(args_lr_mult)\n\n    def set_wd_mult(self, args_wd_mult):\n        self._optimizer.set_wd_mult(args_wd_mult)\n\n\ndef broadcast_parameters(params, root_rank=0):\n    \"\"\"\n    Broadcasts the parameters from root rank to all other processes.\n    Typical usage is to broadcast the `Module.get_params()` or the\n    `Block.collect_params()`.\n\n    Arguments:\n        params: One of the following:\n            - dict of parameters to broadcast\n            - ParameterDict to broadcast\n        root_rank: The rank of the process from which parameters will be\n                   broadcasted to all other processes.\n    \"\"\"\n    tensors = []\n    if isinstance(params, dict):\n        tensors = [p for _, p in sorted(params.items())]\n    elif isinstance(params, mx.gluon.parameter.ParameterDict):\n        for _, p in sorted(params.items()):\n            try:\n                tensors.append(p.data())\n            except mx.gluon.parameter.DeferredInitializationError:\n                # skip broadcasting deferred init param\n                pass\n    else:\n        raise ValueError('invalid params of type: %s' % type(params))\n\n    # Run broadcasts.\n    for i, tensor in enumerate(tensors):\n        broadcast_(tensor, root_rank, str(i))\n\n    # Make sure tensors pushed to MXNet engine get processed such that all\n    # workers are synced before starting training.\n    for tensor in tensors:\n        tensor.wait_to_read()\n", "path": "horovod/mxnet/__init__.py"}], "after_files": [{"content": "import argparse\nimport logging\nimport os\nimport zipfile\nimport time\n\nimport mxnet as mx\nimport horovod.mxnet as hvd\nfrom mxnet import autograd, gluon, nd\nfrom mxnet.test_utils import download\n\n# Training settings\nparser = argparse.ArgumentParser(description='MXNet MNIST Example')\n\nparser.add_argument('--batch-size', type=int, default=64,\n                    help='training batch size (default: 64)')\nparser.add_argument('--dtype', type=str, default='float32',\n                    help='training data type (default: float32)')\nparser.add_argument('--epochs', type=int, default=5,\n                    help='number of training epochs (default: 5)')\nparser.add_argument('--lr', type=float, default=0.01,\n                    help='learning rate (default: 0.01)')\nparser.add_argument('--momentum', type=float, default=0.9,\n                    help='SGD momentum (default: 0.9)')\nparser.add_argument('--no-cuda', action='store_true', default=False,\n                    help='disable training on GPU (default: False)')\nargs = parser.parse_args()\n\nif not args.no_cuda:\n    # Disable CUDA if there are no GPUs.\n    if not mx.test_utils.list_gpus():\n        args.no_cuda = True\n\nlogging.basicConfig(level=logging.INFO)\nlogging.info(args)\n\n\n# Function to get mnist iterator given a rank\ndef get_mnist_iterator(rank):\n    data_dir = \"data-%d\" % rank\n    if not os.path.isdir(data_dir):\n        os.makedirs(data_dir)\n    zip_file_path = download('http://data.mxnet.io/mxnet/data/mnist.zip',\n                             dirname=data_dir)\n    with zipfile.ZipFile(zip_file_path) as zf:\n        zf.extractall(data_dir)\n\n    input_shape = (1, 28, 28)\n    batch_size = args.batch_size\n\n    train_iter = mx.io.MNISTIter(\n        image=\"%s/train-images-idx3-ubyte\" % data_dir,\n        label=\"%s/train-labels-idx1-ubyte\" % data_dir,\n        input_shape=input_shape,\n        batch_size=batch_size,\n        shuffle=True,\n        flat=False,\n        num_parts=hvd.size(),\n        part_index=hvd.rank()\n    )\n\n    val_iter = mx.io.MNISTIter(\n        image=\"%s/t10k-images-idx3-ubyte\" % data_dir,\n        label=\"%s/t10k-labels-idx1-ubyte\" % data_dir,\n        input_shape=input_shape,\n        batch_size=batch_size,\n        flat=False,\n    )\n\n    return train_iter, val_iter\n\n\n# Function to define neural network\ndef conv_nets():\n    net = gluon.nn.HybridSequential()\n    with net.name_scope():\n        net.add(gluon.nn.Conv2D(channels=20, kernel_size=5, activation='relu'))\n        net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))\n        net.add(gluon.nn.Conv2D(channels=50, kernel_size=5, activation='relu'))\n        net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))\n        net.add(gluon.nn.Flatten())\n        net.add(gluon.nn.Dense(512, activation=\"relu\"))\n        net.add(gluon.nn.Dense(10))\n    return net\n\n\n# Function to evaluate accuracy for a model\ndef evaluate(model, data_iter, context):\n    data_iter.reset()\n    metric = mx.metric.Accuracy()\n    for _, batch in enumerate(data_iter):\n        data = batch.data[0].as_in_context(context)\n        label = batch.label[0].as_in_context(context)\n        output = model(data.astype(args.dtype, copy=False))\n        metric.update([label], [output])\n\n    return metric.get()\n\n\n# Initialize Horovod\nhvd.init()\n\n# Horovod: pin context to local rank\ncontext = mx.cpu(hvd.local_rank()) if args.no_cuda else mx.gpu(hvd.local_rank())\nnum_workers = hvd.size()\n\n# Load training and validation data\ntrain_data, val_data = get_mnist_iterator(hvd.rank())\n\n# Build model\nmodel = conv_nets()\nmodel.cast(args.dtype)\nmodel.hybridize()\n\n# Create optimizer\noptimizer_params = {'momentum': args.momentum,\n                    'learning_rate': args.lr * hvd.size(),\n                    'rescale_grad': 1.0 / args.batch_size}\nopt = mx.optimizer.create('sgd', **optimizer_params)\n# Horovod: wrap optimizer with DistributedOptimizer\nopt = hvd.DistributedOptimizer(opt)\n\n# Initialize parameters\ninitializer = mx.init.Xavier(rnd_type='gaussian', factor_type=\"in\",\n                             magnitude=2)\nmodel.initialize(initializer, ctx=context)\n\n# Fetch and broadcast parameters\nparams = model.collect_params()\nif params is not None:\n    hvd.broadcast_parameters(params, root_rank=0)\n\n# Create trainer, loss function and train metric\ntrainer = gluon.Trainer(params, opt, kvstore=None)\nloss_fn = gluon.loss.SoftmaxCrossEntropyLoss()\nmetric = mx.metric.Accuracy()\n\n# Train model\nfor epoch in range(args.epochs):\n    tic = time.time()\n    train_data.reset()\n    metric.reset()\n    for nbatch, batch in enumerate(train_data, start=1):\n        data = batch.data[0].as_in_context(context)\n        label = batch.label[0].as_in_context(context)\n        with autograd.record():\n            output = model(data.astype(args.dtype, copy=False))\n            loss = loss_fn(output, label)\n        loss.backward()\n        trainer.step(args.batch_size)\n        metric.update([label], [output])\n\n        if nbatch % 100 == 0:\n            name, acc = metric.get()\n            logging.info('[Epoch %d Batch %d] Training: %s=%f' %\n                         (epoch, nbatch, name, acc))\n\n    if hvd.rank() == 0:\n        elapsed = time.time() - tic\n        speed = nbatch * args.batch_size * hvd.size() / elapsed\n        logging.info('Epoch[%d]\\tSpeed=%.2f samples/s\\tTime cost=%f',\n                     epoch, speed, elapsed)\n\n    # Evaluate model accuracy\n    _, train_acc = metric.get()\n    name, val_acc = evaluate(model, val_data, context)\n    if hvd.rank() == 0:\n        logging.info('Epoch[%d]\\tTrain: %s=%f\\tValidation: %s=%f', epoch, name,\n                     train_acc, name, val_acc)\n\n    if hvd.rank() == 0 and epoch == args.epochs - 1:\n        assert val_acc > 0.96, \"Achieved accuracy (%f) is lower than expected\\\n                                (0.96)\" % val_acc\n", "path": "examples/mxnet_mnist.py"}, {"content": "# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom horovod.common import check_extension\n\ncheck_extension('horovod.mxnet', 'HOROVOD_WITH_MXNET',\n                __file__, 'mpi_lib')\n\nfrom horovod.mxnet.mpi_ops import allgather\nfrom horovod.mxnet.mpi_ops import allreduce, allreduce_\nfrom horovod.mxnet.mpi_ops import broadcast, broadcast_\nfrom horovod.mxnet.mpi_ops import init, shutdown\nfrom horovod.mxnet.mpi_ops import size, local_size, rank, local_rank\nfrom horovod.mxnet.mpi_ops import mpi_threads_supported\n\nimport mxnet as mx\nimport types\n\n\n# This is where Horovod's DistributedOptimizer wrapper for MXNet goes\nclass DistributedOptimizer(mx.optimizer.Optimizer):\n    def __init__(self, optimizer):\n        self._optimizer = optimizer\n\n    def __getattr__(self, item):\n        return getattr(self._optimizer, item)\n\n    def create_state_multi_precision(self, index, weight):\n        return self._optimizer.create_state_multi_precision(index, weight)\n\n    def _do_allreduce(self, index, grad):\n        if isinstance(index, (tuple, list)):\n            for i in range(len(index)):\n                allreduce_(grad[i], average=True, name=str(index[i]))\n        else:\n            allreduce_(grad, average=True, name=str(index))\n\n    def update(self, index, weight, grad, state):\n        self._do_allreduce(index, grad)\n        self._optimizer.update(index, weight, grad, state)\n\n    def update_multi_precision(self, index, weight, grad, state):\n        self._do_allreduce(index, grad)\n        self._optimizer.update_multi_precision(index, weight, grad, state)\n\n    def set_learning_rate(self, lr):\n        self._optimizer.set_learning_rate(lr)\n\n    def set_lr_mult(self, args_lr_mult):\n        self._optimizer.set_lr_mult(args_lr_mult)\n\n    def set_wd_mult(self, args_wd_mult):\n        self._optimizer.set_wd_mult(args_wd_mult)\n\n\n# Wrapper to inject Horovod broadcast after parameter initialization\ndef _append_broadcast_init(param, root_rank):\n    init_impl = getattr(param, '_init_impl')\n    def wrapped_init_impl(self, *args, **kwargs):\n        init_impl(*args, **kwargs)\n        broadcast_(self.data(), root_rank=root_rank)\n        self.data().wait_to_read()\n    return wrapped_init_impl\n\n\ndef broadcast_parameters(params, root_rank=0):\n    \"\"\"\n    Broadcasts the parameters from root rank to all other processes.\n    Typical usage is to broadcast the `Module.get_params()` or the\n    `Block.collect_params()`.\n\n    Arguments:\n        params: One of the following:\n            - dict of parameters to broadcast\n            - ParameterDict to broadcast\n        root_rank: The rank of the process from which parameters will be\n                   broadcasted to all other processes.\n    \"\"\"\n    tensors = []\n    if isinstance(params, dict):\n        tensors = [p for _, p in sorted(params.items())]\n    elif isinstance(params, mx.gluon.parameter.ParameterDict):\n        for _, p in sorted(params.items()):\n            try:\n                tensors.append(p.data())\n            except mx.gluon.parameter.DeferredInitializationError:\n                # Inject wrapper method with post-initialization broadcast to\n                # handle parameters with deferred initialization\n                new_init = _append_broadcast_init(p, root_rank)\n                p._init_impl = types.MethodType(new_init, p)\n    else:\n        raise ValueError('invalid params of type: %s' % type(params))\n\n    # Run broadcasts.\n    for i, tensor in enumerate(tensors):\n        broadcast_(tensor, root_rank, str(i))\n\n    # Make sure tensors pushed to MXNet engine get processed such that all\n    # workers are synced before starting training.\n    for tensor in tensors:\n        tensor.wait_to_read()\n", "path": "horovod/mxnet/__init__.py"}]}