We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add option to replace bn by sync bn in auto_model
## 🚀 Feature

Idea is to add an option `sync_bn` (default, False) to apply [`convert_sync_batchnorm`](https://pytorch.org/docs/master/generated/torch.nn.SyncBatchNorm.html#torch.nn.SyncBatchNorm.convert_sync_batchnorm) on the input model if a DDP wrapper is used further.

- Check performances on PascalVOC
- Update ref examples


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `ignite/distributed/auto.py`
Content:
```
1 import warnings
2 
3 import torch
4 import torch.nn as nn
5 from torch.optim.optimizer import Optimizer
6 from torch.utils.data import DataLoader, Dataset
7 from torch.utils.data.distributed import DistributedSampler
8 from torch.utils.data.sampler import Sampler
9 
10 from ignite.distributed import utils as idist
11 from ignite.distributed.comp_models import horovod as idist_hvd
12 from ignite.distributed.comp_models import native as idist_native
13 from ignite.distributed.comp_models import xla as idist_xla
14 from ignite.utils import setup_logger
15 
16 __all__ = ["auto_dataloader", "auto_model", "auto_optim", "DistributedProxySampler"]
17 
18 
19 def auto_dataloader(dataset, **kwargs):
20     """Helper method to create a dataloader adapted for non-distributed and distributed configurations (supporting
21     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).
22 
23     Internally, we create a dataloader with provided kwargs while applying the following updates:
24 
25     - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.
26     - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.
27     - if no sampler provided by user, `torch DistributedSampler` is setup.
28     - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.
29     - if the default device is 'cuda', `pin_memory` is automatically set to `True`.
30 
31     .. warning::
32 
33         Custom batch sampler is not adapted for distributed configuration. Please, make sure that provided batch
34         sampler is compatible with distributed configuration.
35 
36     Examples:
37 
38     .. code-block:: python
39 
40         import ignite.distribted as idist
41 
42         train_loader = idist.auto_dataloader(
43             train_dataset,
44             batch_size=32,
45             num_workers=4,
46             shuffle=True,
47             pin_memory="cuda" in idist.device().type,
48             drop_last=True,
49         )
50 
51     Args:
52         dataset (Dataset): input torch dataset
53         **kwargs: keyword arguments for `torch DataLoader`_.
54 
55     Returns:
56         `torch DataLoader`_ or `XLA MpDeviceLoader`_ for XLA devices
57 
58     .. _torch DataLoader: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
59     .. _XLA MpDeviceLoader: https://github.com/pytorch/xla/blob/master/torch_xla/distributed/parallel_loader.py#L178
60     .. _torch DistributedSampler:
61         https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
62     """
63     rank = idist.get_rank()
64     world_size = idist.get_world_size()
65 
66     logger = setup_logger(__name__ + ".auto_dataloader")
67     if world_size > 1:
68         if "batch_size" in kwargs and kwargs["batch_size"] >= world_size:
69             kwargs["batch_size"] //= world_size
70 
71         nproc = idist.get_nproc_per_node()
72         if "num_workers" in kwargs and kwargs["num_workers"] >= nproc:
73             kwargs["num_workers"] = (kwargs["num_workers"] + nproc - 1) // nproc
74 
75         if "batch_sampler" not in kwargs:
76             if kwargs.get("sampler", None) is not None:
77                 sampler = DistributedProxySampler(kwargs["sampler"], num_replicas=world_size, rank=rank)
78             else:
79                 sampler = DistributedSampler(
80                     dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get("shuffle", True)
81                 )
82                 # we need to remove "shuffle" from kwargs if sampler is used
83                 if "shuffle" in kwargs:
84                     del kwargs["shuffle"]
85 
86             kwargs["sampler"] = sampler
87         else:
88             warnings.warn(
89                 "Found batch_sampler in provided kwargs. Please, make sure that it is compatible "
90                 "with distributed configuration"
91             )
92 
93     if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and kwargs.get("pin_memory", False):
94         # TODO: How about XLA GPU ?
95         warnings.warn(
96             "Found incompatible options: xla support and pin_memory args equal True. "
97             "Argument `pin_memory=False` will be used to construct data loader."
98         )
99         kwargs["pin_memory"] = False
100     else:
101         kwargs["pin_memory"] = kwargs.get("pin_memory", "cuda" in idist.device().type)
102 
103     logger.info("Use data loader kwargs for dataset '{}': \n\t{}".format(repr(dataset)[:20].strip(), kwargs))
104     dataloader = DataLoader(dataset, **kwargs)
105 
106     if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and world_size > 1:
107 
108         logger.info("DataLoader is wrapped by `MpDeviceLoader` on XLA")
109 
110         mp_device_loader_cls = _MpDeviceLoader
111         try:
112             from torch_xla.distributed.parallel_loader import MpDeviceLoader
113 
114             mp_device_loader_cls = MpDeviceLoader
115         except ImportError:
116             pass
117 
118         sampler = dataloader.sampler
119         dataloader = mp_device_loader_cls(dataloader, idist.device())
120         dataloader.sampler = sampler
121 
122     return dataloader
123 
124 
125 def auto_model(model: nn.Module) -> nn.Module:
126     """Helper method to adapt provided model for non-distributed and distributed configurations (supporting
127     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).
128 
129     Internally, we perform to following:
130 
131     - send model to current :meth:`~ignite.distributed.utils.device()` if model's parameters are not on the device.
132     - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1.
133     - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.
134     - broadcast the initial variable states from rank 0 to all other processes if Horovod distributed framework is used.
135 
136     Examples:
137 
138     .. code-block:: python
139 
140         import ignite.distribted as idist
141 
142         model = idist.auto_model(model)
143 
144     In addition with NVidia/Apex, it can be used in the following way:
145 
146     .. code-block:: python
147 
148         import ignite.distribted as idist
149 
150         model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
151         model = idist.auto_model(model)
152 
153     Args:
154         model (torch.nn.Module): model to adapt.
155 
156     Returns:
157         torch.nn.Module
158 
159     .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel
160     .. _torch DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel
161     """
162     logger = setup_logger(__name__ + ".auto_model")
163 
164     # Put model's parameters to device if its parameters are not on the device
165     device = idist.device()
166     if not all([p.device == device for p in model.parameters()]):
167         model.to(device)
168 
169     # distributed data parallel model
170     if idist.get_world_size() > 1:
171         bnd = idist.backend()
172         if idist.has_native_dist_support and bnd == idist_native.NCCL:
173             lrank = idist.get_local_rank()
174             logger.info("Apply torch DistributedDataParallel on model, device id: {}".format(lrank))
175             model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])
176         elif idist.has_native_dist_support and bnd == idist_native.GLOO:
177             logger.info("Apply torch DistributedDataParallel on model")
178             model = torch.nn.parallel.DistributedDataParallel(model)
179         elif idist.has_hvd_support and bnd == idist_hvd.HOROVOD:
180             import horovod.torch as hvd
181 
182             logger.info("Broadcast the initial variable states from rank 0 to all other processes")
183             hvd.broadcast_parameters(model.state_dict(), root_rank=0)
184 
185     # not distributed but multiple GPUs reachable so data parallel model
186     elif torch.cuda.device_count() > 1 and "cuda" in idist.device().type:
187         logger.info("Apply torch DataParallel on model")
188         model = torch.nn.parallel.DataParallel(model)
189 
190     return model
191 
192 
193 def auto_optim(optimizer: Optimizer) -> Optimizer:
194     """Helper method to adapt optimizer for non-distributed and distributed configurations (supporting
195     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).
196 
197     Internally, this method is no-op for non-distributed and torch native distributed configuration.
198 
199     For XLA distributed configuration, we create a new class that inherits from provided optimizer.
200     The goal is to override the `step()` method with specific `xm.optimizer_step`_ implementation.
201 
202     For Horovod distributed configuration, optimizer is wrapped with Horovod Distributed Optimizer and
203     its state is broadcasted from rank 0 to all other processes.
204 
205     Examples:
206 
207     .. code-block:: python
208 
209         import ignite.distributed as idist
210 
211         optimizer = idist.auto_optim(optimizer)
212 
213 
214     Args:
215         optimizer (Optimizer): input torch optimizer
216 
217     Returns:
218         Optimizer
219 
220     .. _xm.optimizer_step: http://pytorch.org/xla/release/1.5/index.html#torch_xla.core.xla_model.optimizer_step
221 
222     """
223     bnd = idist.backend()
224     if idist.has_xla_support and bnd == idist_xla.XLA_TPU:
225         cls = type(optimizer.__class__.__name__, (optimizer.__class__,), dict(_XLADistributedOptimizer.__dict__))
226         return cls(optimizer)
227 
228     if idist.has_hvd_support and bnd == idist_hvd.HOROVOD:
229         import horovod.torch as hvd
230 
231         optimizer = hvd.DistributedOptimizer(optimizer)
232         hvd.broadcast_optimizer_state(optimizer, root_rank=0)
233         return optimizer
234 
235     return optimizer
236 
237 
238 class DistributedProxySampler(DistributedSampler):
239     """Distributed sampler proxy to adapt user's sampler for distributed data parallelism configuration.
240 
241     Code is based on https://github.com/pytorch/pytorch/issues/23430#issuecomment-562350407
242 
243 
244     .. note::
245         Input sampler is assumed to have a constant size.
246 
247     Args:
248         sampler (Sampler): Input torch data sampler.
249         num_replicas (int, optional): Number of processes participating in distributed training.
250         rank (int, optional): Rank of the current process within ``num_replicas``.
251 
252     """
253 
254     def __init__(self, sampler: Sampler, num_replicas=None, rank=None):
255 
256         if not isinstance(sampler, Sampler):
257             raise TypeError("Argument sampler should be instance of torch Sampler, but given: {}".format(type(sampler)))
258 
259         if not hasattr(sampler, "__len__"):
260             raise TypeError("Argument sampler should have length")
261 
262         super(DistributedProxySampler, self).__init__(sampler, num_replicas=num_replicas, rank=rank, shuffle=False)
263         self.sampler = sampler
264 
265     def __iter__(self):
266         # deterministically shuffle based on epoch
267         torch.manual_seed(self.epoch)
268 
269         indices = []
270         while len(indices) < self.total_size:
271             indices += list(self.sampler)
272 
273         if len(indices) > self.total_size:
274             indices = indices[: self.total_size]
275 
276         # subsample
277         indices = indices[self.rank : self.total_size : self.num_replicas]
278         if len(indices) != self.num_samples:
279             raise RuntimeError("{} vs {}".format(len(indices), self.num_samples))
280 
281         return iter(indices)
282 
283 
284 if idist.has_xla_support:
285 
286     import torch_xla.core.xla_model as xm
287     from torch_xla.distributed.parallel_loader import ParallelLoader
288 
289     class _MpDeviceLoader:
290         # https://github.com/pytorch/xla/pull/2117
291         # From pytorch/xla if `torch_xla.distributed.parallel_loader.MpDeviceLoader` is not available
292         def __init__(self, loader, device, **kwargs):
293             self._loader = loader
294             self._device = device
295             self._parallel_loader_kwargs = kwargs
296 
297         def __iter__(self):
298             parallel_loader = ParallelLoader(self._loader, [self._device], **self._parallel_loader_kwargs)
299             return parallel_loader.per_device_loader(self._device)
300 
301         def __len__(self):
302             return len(self._loader)
303 
304     class _XLADistributedOptimizer(Optimizer):
305         def __init__(self, optimizer):
306             super(self.__class__, self).__init__(optimizer.param_groups)
307             self.wrapped_optimizer = optimizer
308 
309         def step(self, closure=None):
310             xm.optimizer_step(self.wrapped_optimizer, barrier=True)
311
```
--- 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/ignite/distributed/auto.py b/ignite/distributed/auto.py
--- a/ignite/distributed/auto.py
+++ b/ignite/distributed/auto.py
@@ -24,7 +24,7 @@
 
     - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.
     - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.
-    - if no sampler provided by user, `torch DistributedSampler` is setup.
+    - if no sampler provided by user, `torch DistributedSampler`_ is setup.
     - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.
     - if the default device is 'cuda', `pin_memory` is automatically set to `True`.
 
@@ -122,7 +122,7 @@
     return dataloader
 
 
-def auto_model(model: nn.Module) -> nn.Module:
+def auto_model(model: nn.Module, sync_bn: bool = False) -> nn.Module:
     """Helper method to adapt provided model for non-distributed and distributed configurations (supporting
     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).
 
@@ -152,12 +152,19 @@
 
     Args:
         model (torch.nn.Module): model to adapt.
+        sync_bn (bool): if True, applies `torch convert_sync_batchnorm`_ to the model for native torch
+            distributed only. Default, False. Note, if using Nvidia/Apex, batchnorm conversion should be
+            applied before calling ``amp.initialize``.
 
     Returns:
         torch.nn.Module
 
-    .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel
-    .. _torch DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel
+    .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/generated/torch.nn.parallel.
+        DistributedDataParallel.html
+    .. _torch DataParallel: https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html
+    .. _torch convert_sync_batchnorm: https://pytorch.org/docs/stable/generated/torch.nn.SyncBatchNorm.html#
+        torch.nn.SyncBatchNorm.convert_sync_batchnorm
+
     """
     logger = setup_logger(__name__ + ".auto_model")
 
@@ -170,10 +177,18 @@
     if idist.get_world_size() > 1:
         bnd = idist.backend()
         if idist.has_native_dist_support and bnd == idist_native.NCCL:
+            if sync_bn:
+                logger.info("Convert batch norm to sync batch norm")
+                model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
             lrank = idist.get_local_rank()
             logger.info("Apply torch DistributedDataParallel on model, device id: {}".format(lrank))
             model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])
         elif idist.has_native_dist_support and bnd == idist_native.GLOO:
+            if sync_bn:
+                logger.info("Convert batch norm to sync batch norm")
+                model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
             logger.info("Apply torch DistributedDataParallel on model")
             model = torch.nn.parallel.DistributedDataParallel(model)
         elif idist.has_hvd_support and bnd == idist_hvd.HOROVOD:

{"golden_diff": "diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py\n--- a/ignite/distributed/auto.py\n+++ b/ignite/distributed/auto.py\n@@ -24,7 +24,7 @@\n \n     - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.\n     - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.\n-    - if no sampler provided by user, `torch DistributedSampler` is setup.\n+    - if no sampler provided by user, `torch DistributedSampler`_ is setup.\n     - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.\n     - if the default device is 'cuda', `pin_memory` is automatically set to `True`.\n \n@@ -122,7 +122,7 @@\n     return dataloader\n \n \n-def auto_model(model: nn.Module) -> nn.Module:\n+def auto_model(model: nn.Module, sync_bn: bool = False) -> nn.Module:\n     \"\"\"Helper method to adapt provided model for non-distributed and distributed configurations (supporting\n     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n \n@@ -152,12 +152,19 @@\n \n     Args:\n         model (torch.nn.Module): model to adapt.\n+        sync_bn (bool): if True, applies `torch convert_sync_batchnorm`_ to the model for native torch\n+            distributed only. Default, False. Note, if using Nvidia/Apex, batchnorm conversion should be\n+            applied before calling ``amp.initialize``.\n \n     Returns:\n         torch.nn.Module\n \n-    .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel\n-    .. _torch DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel\n+    .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/generated/torch.nn.parallel.\n+        DistributedDataParallel.html\n+    .. _torch DataParallel: https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html\n+    .. _torch convert_sync_batchnorm: https://pytorch.org/docs/stable/generated/torch.nn.SyncBatchNorm.html#\n+        torch.nn.SyncBatchNorm.convert_sync_batchnorm\n+\n     \"\"\"\n     logger = setup_logger(__name__ + \".auto_model\")\n \n@@ -170,10 +177,18 @@\n     if idist.get_world_size() > 1:\n         bnd = idist.backend()\n         if idist.has_native_dist_support and bnd == idist_native.NCCL:\n+            if sync_bn:\n+                logger.info(\"Convert batch norm to sync batch norm\")\n+                model = nn.SyncBatchNorm.convert_sync_batchnorm(model)\n+\n             lrank = idist.get_local_rank()\n             logger.info(\"Apply torch DistributedDataParallel on model, device id: {}\".format(lrank))\n             model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])\n         elif idist.has_native_dist_support and bnd == idist_native.GLOO:\n+            if sync_bn:\n+                logger.info(\"Convert batch norm to sync batch norm\")\n+                model = nn.SyncBatchNorm.convert_sync_batchnorm(model)\n+\n             logger.info(\"Apply torch DistributedDataParallel on model\")\n             model = torch.nn.parallel.DistributedDataParallel(model)\n         elif idist.has_hvd_support and bnd == idist_hvd.HOROVOD:\n", "issue": "Add option to replace bn by sync bn in auto_model\n## \ud83d\ude80 Feature\r\n\r\nIdea is to add an option `sync_bn` (default, False) to apply [`convert_sync_batchnorm`](https://pytorch.org/docs/master/generated/torch.nn.SyncBatchNorm.html#torch.nn.SyncBatchNorm.convert_sync_batchnorm) on the input model if a DDP wrapper is used further.\r\n\r\n- Check performances on PascalVOC\r\n- Update ref examples\r\n\n", "before_files": [{"content": "import warnings\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim.optimizer import Optimizer\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.data.sampler import Sampler\n\nfrom ignite.distributed import utils as idist\nfrom ignite.distributed.comp_models import horovod as idist_hvd\nfrom ignite.distributed.comp_models import native as idist_native\nfrom ignite.distributed.comp_models import xla as idist_xla\nfrom ignite.utils import setup_logger\n\n__all__ = [\"auto_dataloader\", \"auto_model\", \"auto_optim\", \"DistributedProxySampler\"]\n\n\ndef auto_dataloader(dataset, **kwargs):\n    \"\"\"Helper method to create a dataloader adapted for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, we create a dataloader with provided kwargs while applying the following updates:\n\n    - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.\n    - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.\n    - if no sampler provided by user, `torch DistributedSampler` is setup.\n    - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.\n    - if the default device is 'cuda', `pin_memory` is automatically set to `True`.\n\n    .. warning::\n\n        Custom batch sampler is not adapted for distributed configuration. Please, make sure that provided batch\n        sampler is compatible with distributed configuration.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        train_loader = idist.auto_dataloader(\n            train_dataset,\n            batch_size=32,\n            num_workers=4,\n            shuffle=True,\n            pin_memory=\"cuda\" in idist.device().type,\n            drop_last=True,\n        )\n\n    Args:\n        dataset (Dataset): input torch dataset\n        **kwargs: keyword arguments for `torch DataLoader`_.\n\n    Returns:\n        `torch DataLoader`_ or `XLA MpDeviceLoader`_ for XLA devices\n\n    .. _torch DataLoader: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader\n    .. _XLA MpDeviceLoader: https://github.com/pytorch/xla/blob/master/torch_xla/distributed/parallel_loader.py#L178\n    .. _torch DistributedSampler:\n        https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler\n    \"\"\"\n    rank = idist.get_rank()\n    world_size = idist.get_world_size()\n\n    logger = setup_logger(__name__ + \".auto_dataloader\")\n    if world_size > 1:\n        if \"batch_size\" in kwargs and kwargs[\"batch_size\"] >= world_size:\n            kwargs[\"batch_size\"] //= world_size\n\n        nproc = idist.get_nproc_per_node()\n        if \"num_workers\" in kwargs and kwargs[\"num_workers\"] >= nproc:\n            kwargs[\"num_workers\"] = (kwargs[\"num_workers\"] + nproc - 1) // nproc\n\n        if \"batch_sampler\" not in kwargs:\n            if kwargs.get(\"sampler\", None) is not None:\n                sampler = DistributedProxySampler(kwargs[\"sampler\"], num_replicas=world_size, rank=rank)\n            else:\n                sampler = DistributedSampler(\n                    dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get(\"shuffle\", True)\n                )\n                # we need to remove \"shuffle\" from kwargs if sampler is used\n                if \"shuffle\" in kwargs:\n                    del kwargs[\"shuffle\"]\n\n            kwargs[\"sampler\"] = sampler\n        else:\n            warnings.warn(\n                \"Found batch_sampler in provided kwargs. Please, make sure that it is compatible \"\n                \"with distributed configuration\"\n            )\n\n    if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and kwargs.get(\"pin_memory\", False):\n        # TODO: How about XLA GPU ?\n        warnings.warn(\n            \"Found incompatible options: xla support and pin_memory args equal True. \"\n            \"Argument `pin_memory=False` will be used to construct data loader.\"\n        )\n        kwargs[\"pin_memory\"] = False\n    else:\n        kwargs[\"pin_memory\"] = kwargs.get(\"pin_memory\", \"cuda\" in idist.device().type)\n\n    logger.info(\"Use data loader kwargs for dataset '{}': \\n\\t{}\".format(repr(dataset)[:20].strip(), kwargs))\n    dataloader = DataLoader(dataset, **kwargs)\n\n    if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and world_size > 1:\n\n        logger.info(\"DataLoader is wrapped by `MpDeviceLoader` on XLA\")\n\n        mp_device_loader_cls = _MpDeviceLoader\n        try:\n            from torch_xla.distributed.parallel_loader import MpDeviceLoader\n\n            mp_device_loader_cls = MpDeviceLoader\n        except ImportError:\n            pass\n\n        sampler = dataloader.sampler\n        dataloader = mp_device_loader_cls(dataloader, idist.device())\n        dataloader.sampler = sampler\n\n    return dataloader\n\n\ndef auto_model(model: nn.Module) -> nn.Module:\n    \"\"\"Helper method to adapt provided model for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, we perform to following:\n\n    - send model to current :meth:`~ignite.distributed.utils.device()` if model's parameters are not on the device.\n    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1.\n    - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.\n    - broadcast the initial variable states from rank 0 to all other processes if Horovod distributed framework is used.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        model = idist.auto_model(model)\n\n    In addition with NVidia/Apex, it can be used in the following way:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)\n        model = idist.auto_model(model)\n\n    Args:\n        model (torch.nn.Module): model to adapt.\n\n    Returns:\n        torch.nn.Module\n\n    .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel\n    .. _torch DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel\n    \"\"\"\n    logger = setup_logger(__name__ + \".auto_model\")\n\n    # Put model's parameters to device if its parameters are not on the device\n    device = idist.device()\n    if not all([p.device == device for p in model.parameters()]):\n        model.to(device)\n\n    # distributed data parallel model\n    if idist.get_world_size() > 1:\n        bnd = idist.backend()\n        if idist.has_native_dist_support and bnd == idist_native.NCCL:\n            lrank = idist.get_local_rank()\n            logger.info(\"Apply torch DistributedDataParallel on model, device id: {}\".format(lrank))\n            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])\n        elif idist.has_native_dist_support and bnd == idist_native.GLOO:\n            logger.info(\"Apply torch DistributedDataParallel on model\")\n            model = torch.nn.parallel.DistributedDataParallel(model)\n        elif idist.has_hvd_support and bnd == idist_hvd.HOROVOD:\n            import horovod.torch as hvd\n\n            logger.info(\"Broadcast the initial variable states from rank 0 to all other processes\")\n            hvd.broadcast_parameters(model.state_dict(), root_rank=0)\n\n    # not distributed but multiple GPUs reachable so data parallel model\n    elif torch.cuda.device_count() > 1 and \"cuda\" in idist.device().type:\n        logger.info(\"Apply torch DataParallel on model\")\n        model = torch.nn.parallel.DataParallel(model)\n\n    return model\n\n\ndef auto_optim(optimizer: Optimizer) -> Optimizer:\n    \"\"\"Helper method to adapt optimizer for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, this method is no-op for non-distributed and torch native distributed configuration.\n\n    For XLA distributed configuration, we create a new class that inherits from provided optimizer.\n    The goal is to override the `step()` method with specific `xm.optimizer_step`_ implementation.\n\n    For Horovod distributed configuration, optimizer is wrapped with Horovod Distributed Optimizer and\n    its state is broadcasted from rank 0 to all other processes.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distributed as idist\n\n        optimizer = idist.auto_optim(optimizer)\n\n\n    Args:\n        optimizer (Optimizer): input torch optimizer\n\n    Returns:\n        Optimizer\n\n    .. _xm.optimizer_step: http://pytorch.org/xla/release/1.5/index.html#torch_xla.core.xla_model.optimizer_step\n\n    \"\"\"\n    bnd = idist.backend()\n    if idist.has_xla_support and bnd == idist_xla.XLA_TPU:\n        cls = type(optimizer.__class__.__name__, (optimizer.__class__,), dict(_XLADistributedOptimizer.__dict__))\n        return cls(optimizer)\n\n    if idist.has_hvd_support and bnd == idist_hvd.HOROVOD:\n        import horovod.torch as hvd\n\n        optimizer = hvd.DistributedOptimizer(optimizer)\n        hvd.broadcast_optimizer_state(optimizer, root_rank=0)\n        return optimizer\n\n    return optimizer\n\n\nclass DistributedProxySampler(DistributedSampler):\n    \"\"\"Distributed sampler proxy to adapt user's sampler for distributed data parallelism configuration.\n\n    Code is based on https://github.com/pytorch/pytorch/issues/23430#issuecomment-562350407\n\n\n    .. note::\n        Input sampler is assumed to have a constant size.\n\n    Args:\n        sampler (Sampler): Input torch data sampler.\n        num_replicas (int, optional): Number of processes participating in distributed training.\n        rank (int, optional): Rank of the current process within ``num_replicas``.\n\n    \"\"\"\n\n    def __init__(self, sampler: Sampler, num_replicas=None, rank=None):\n\n        if not isinstance(sampler, Sampler):\n            raise TypeError(\"Argument sampler should be instance of torch Sampler, but given: {}\".format(type(sampler)))\n\n        if not hasattr(sampler, \"__len__\"):\n            raise TypeError(\"Argument sampler should have length\")\n\n        super(DistributedProxySampler, self).__init__(sampler, num_replicas=num_replicas, rank=rank, shuffle=False)\n        self.sampler = sampler\n\n    def __iter__(self):\n        # deterministically shuffle based on epoch\n        torch.manual_seed(self.epoch)\n\n        indices = []\n        while len(indices) < self.total_size:\n            indices += list(self.sampler)\n\n        if len(indices) > self.total_size:\n            indices = indices[: self.total_size]\n\n        # subsample\n        indices = indices[self.rank : self.total_size : self.num_replicas]\n        if len(indices) != self.num_samples:\n            raise RuntimeError(\"{} vs {}\".format(len(indices), self.num_samples))\n\n        return iter(indices)\n\n\nif idist.has_xla_support:\n\n    import torch_xla.core.xla_model as xm\n    from torch_xla.distributed.parallel_loader import ParallelLoader\n\n    class _MpDeviceLoader:\n        # https://github.com/pytorch/xla/pull/2117\n        # From pytorch/xla if `torch_xla.distributed.parallel_loader.MpDeviceLoader` is not available\n        def __init__(self, loader, device, **kwargs):\n            self._loader = loader\n            self._device = device\n            self._parallel_loader_kwargs = kwargs\n\n        def __iter__(self):\n            parallel_loader = ParallelLoader(self._loader, [self._device], **self._parallel_loader_kwargs)\n            return parallel_loader.per_device_loader(self._device)\n\n        def __len__(self):\n            return len(self._loader)\n\n    class _XLADistributedOptimizer(Optimizer):\n        def __init__(self, optimizer):\n            super(self.__class__, self).__init__(optimizer.param_groups)\n            self.wrapped_optimizer = optimizer\n\n        def step(self, closure=None):\n            xm.optimizer_step(self.wrapped_optimizer, barrier=True)\n", "path": "ignite/distributed/auto.py"}], "after_files": [{"content": "import warnings\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim.optimizer import Optimizer\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.data.sampler import Sampler\n\nfrom ignite.distributed import utils as idist\nfrom ignite.distributed.comp_models import horovod as idist_hvd\nfrom ignite.distributed.comp_models import native as idist_native\nfrom ignite.distributed.comp_models import xla as idist_xla\nfrom ignite.utils import setup_logger\n\n__all__ = [\"auto_dataloader\", \"auto_model\", \"auto_optim\", \"DistributedProxySampler\"]\n\n\ndef auto_dataloader(dataset, **kwargs):\n    \"\"\"Helper method to create a dataloader adapted for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, we create a dataloader with provided kwargs while applying the following updates:\n\n    - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.\n    - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.\n    - if no sampler provided by user, `torch DistributedSampler`_ is setup.\n    - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.\n    - if the default device is 'cuda', `pin_memory` is automatically set to `True`.\n\n    .. warning::\n\n        Custom batch sampler is not adapted for distributed configuration. Please, make sure that provided batch\n        sampler is compatible with distributed configuration.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        train_loader = idist.auto_dataloader(\n            train_dataset,\n            batch_size=32,\n            num_workers=4,\n            shuffle=True,\n            pin_memory=\"cuda\" in idist.device().type,\n            drop_last=True,\n        )\n\n    Args:\n        dataset (Dataset): input torch dataset\n        **kwargs: keyword arguments for `torch DataLoader`_.\n\n    Returns:\n        `torch DataLoader`_ or `XLA MpDeviceLoader`_ for XLA devices\n\n    .. _torch DataLoader: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader\n    .. _XLA MpDeviceLoader: https://github.com/pytorch/xla/blob/master/torch_xla/distributed/parallel_loader.py#L178\n    .. _torch DistributedSampler:\n        https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler\n    \"\"\"\n    rank = idist.get_rank()\n    world_size = idist.get_world_size()\n\n    logger = setup_logger(__name__ + \".auto_dataloader\")\n    if world_size > 1:\n        if \"batch_size\" in kwargs and kwargs[\"batch_size\"] >= world_size:\n            kwargs[\"batch_size\"] //= world_size\n\n        nproc = idist.get_nproc_per_node()\n        if \"num_workers\" in kwargs and kwargs[\"num_workers\"] >= nproc:\n            kwargs[\"num_workers\"] = (kwargs[\"num_workers\"] + nproc - 1) // nproc\n\n        if \"batch_sampler\" not in kwargs:\n            if kwargs.get(\"sampler\", None) is not None:\n                sampler = DistributedProxySampler(kwargs[\"sampler\"], num_replicas=world_size, rank=rank)\n            else:\n                sampler = DistributedSampler(\n                    dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get(\"shuffle\", True)\n                )\n                # we need to remove \"shuffle\" from kwargs if sampler is used\n                if \"shuffle\" in kwargs:\n                    del kwargs[\"shuffle\"]\n\n            kwargs[\"sampler\"] = sampler\n        else:\n            warnings.warn(\n                \"Found batch_sampler in provided kwargs. Please, make sure that it is compatible \"\n                \"with distributed configuration\"\n            )\n\n    if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and kwargs.get(\"pin_memory\", False):\n        # TODO: How about XLA GPU ?\n        warnings.warn(\n            \"Found incompatible options: xla support and pin_memory args equal True. \"\n            \"Argument `pin_memory=False` will be used to construct data loader.\"\n        )\n        kwargs[\"pin_memory\"] = False\n    else:\n        kwargs[\"pin_memory\"] = kwargs.get(\"pin_memory\", \"cuda\" in idist.device().type)\n\n    logger.info(\"Use data loader kwargs for dataset '{}': \\n\\t{}\".format(repr(dataset)[:20].strip(), kwargs))\n    dataloader = DataLoader(dataset, **kwargs)\n\n    if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and world_size > 1:\n\n        logger.info(\"DataLoader is wrapped by `MpDeviceLoader` on XLA\")\n\n        mp_device_loader_cls = _MpDeviceLoader\n        try:\n            from torch_xla.distributed.parallel_loader import MpDeviceLoader\n\n            mp_device_loader_cls = MpDeviceLoader\n        except ImportError:\n            pass\n\n        sampler = dataloader.sampler\n        dataloader = mp_device_loader_cls(dataloader, idist.device())\n        dataloader.sampler = sampler\n\n    return dataloader\n\n\ndef auto_model(model: nn.Module, sync_bn: bool = False) -> nn.Module:\n    \"\"\"Helper method to adapt provided model for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, we perform to following:\n\n    - send model to current :meth:`~ignite.distributed.utils.device()` if model's parameters are not on the device.\n    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1.\n    - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.\n    - broadcast the initial variable states from rank 0 to all other processes if Horovod distributed framework is used.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        model = idist.auto_model(model)\n\n    In addition with NVidia/Apex, it can be used in the following way:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)\n        model = idist.auto_model(model)\n\n    Args:\n        model (torch.nn.Module): model to adapt.\n        sync_bn (bool): if True, applies `torch convert_sync_batchnorm`_ to the model for native torch\n            distributed only. Default, False. Note, if using Nvidia/Apex, batchnorm conversion should be\n            applied before calling ``amp.initialize``.\n\n    Returns:\n        torch.nn.Module\n\n    .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/generated/torch.nn.parallel.\n        DistributedDataParallel.html\n    .. _torch DataParallel: https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html\n    .. _torch convert_sync_batchnorm: https://pytorch.org/docs/stable/generated/torch.nn.SyncBatchNorm.html#\n        torch.nn.SyncBatchNorm.convert_sync_batchnorm\n\n    \"\"\"\n    logger = setup_logger(__name__ + \".auto_model\")\n\n    # Put model's parameters to device if its parameters are not on the device\n    device = idist.device()\n    if not all([p.device == device for p in model.parameters()]):\n        model.to(device)\n\n    # distributed data parallel model\n    if idist.get_world_size() > 1:\n        bnd = idist.backend()\n        if idist.has_native_dist_support and bnd == idist_native.NCCL:\n            if sync_bn:\n                logger.info(\"Convert batch norm to sync batch norm\")\n                model = nn.SyncBatchNorm.convert_sync_batchnorm(model)\n\n            lrank = idist.get_local_rank()\n            logger.info(\"Apply torch DistributedDataParallel on model, device id: {}\".format(lrank))\n            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])\n        elif idist.has_native_dist_support and bnd == idist_native.GLOO:\n            if sync_bn:\n                logger.info(\"Convert batch norm to sync batch norm\")\n                model = nn.SyncBatchNorm.convert_sync_batchnorm(model)\n\n            logger.info(\"Apply torch DistributedDataParallel on model\")\n            model = torch.nn.parallel.DistributedDataParallel(model)\n        elif idist.has_hvd_support and bnd == idist_hvd.HOROVOD:\n            import horovod.torch as hvd\n\n            logger.info(\"Broadcast the initial variable states from rank 0 to all other processes\")\n            hvd.broadcast_parameters(model.state_dict(), root_rank=0)\n\n    # not distributed but multiple GPUs reachable so data parallel model\n    elif torch.cuda.device_count() > 1 and \"cuda\" in idist.device().type:\n        logger.info(\"Apply torch DataParallel on model\")\n        model = torch.nn.parallel.DataParallel(model)\n\n    return model\n\n\ndef auto_optim(optimizer: Optimizer) -> Optimizer:\n    \"\"\"Helper method to adapt optimizer for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, this method is no-op for non-distributed and torch native distributed configuration.\n\n    For XLA distributed configuration, we create a new class that inherits from provided optimizer.\n    The goal is to override the `step()` method with specific `xm.optimizer_step`_ implementation.\n\n    For Horovod distributed configuration, optimizer is wrapped with Horovod Distributed Optimizer and\n    its state is broadcasted from rank 0 to all other processes.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distributed as idist\n\n        optimizer = idist.auto_optim(optimizer)\n\n\n    Args:\n        optimizer (Optimizer): input torch optimizer\n\n    Returns:\n        Optimizer\n\n    .. _xm.optimizer_step: http://pytorch.org/xla/release/1.5/index.html#torch_xla.core.xla_model.optimizer_step\n\n    \"\"\"\n    bnd = idist.backend()\n    if idist.has_xla_support and bnd == idist_xla.XLA_TPU:\n        cls = type(optimizer.__class__.__name__, (optimizer.__class__,), dict(_XLADistributedOptimizer.__dict__))\n        return cls(optimizer)\n\n    if idist.has_hvd_support and bnd == idist_hvd.HOROVOD:\n        import horovod.torch as hvd\n\n        optimizer = hvd.DistributedOptimizer(optimizer)\n        hvd.broadcast_optimizer_state(optimizer, root_rank=0)\n        return optimizer\n\n    return optimizer\n\n\nclass DistributedProxySampler(DistributedSampler):\n    \"\"\"Distributed sampler proxy to adapt user's sampler for distributed data parallelism configuration.\n\n    Code is based on https://github.com/pytorch/pytorch/issues/23430#issuecomment-562350407\n\n\n    .. note::\n        Input sampler is assumed to have a constant size.\n\n    Args:\n        sampler (Sampler): Input torch data sampler.\n        num_replicas (int, optional): Number of processes participating in distributed training.\n        rank (int, optional): Rank of the current process within ``num_replicas``.\n\n    \"\"\"\n\n    def __init__(self, sampler: Sampler, num_replicas=None, rank=None):\n\n        if not isinstance(sampler, Sampler):\n            raise TypeError(\"Argument sampler should be instance of torch Sampler, but given: {}\".format(type(sampler)))\n\n        if not hasattr(sampler, \"__len__\"):\n            raise TypeError(\"Argument sampler should have length\")\n\n        super(DistributedProxySampler, self).__init__(sampler, num_replicas=num_replicas, rank=rank, shuffle=False)\n        self.sampler = sampler\n\n    def __iter__(self):\n        # deterministically shuffle based on epoch\n        torch.manual_seed(self.epoch)\n\n        indices = []\n        while len(indices) < self.total_size:\n            indices += list(self.sampler)\n\n        if len(indices) > self.total_size:\n            indices = indices[: self.total_size]\n\n        # subsample\n        indices = indices[self.rank : self.total_size : self.num_replicas]\n        if len(indices) != self.num_samples:\n            raise RuntimeError(\"{} vs {}\".format(len(indices), self.num_samples))\n\n        return iter(indices)\n\n\nif idist.has_xla_support:\n\n    import torch_xla.core.xla_model as xm\n    from torch_xla.distributed.parallel_loader import ParallelLoader\n\n    class _MpDeviceLoader:\n        # https://github.com/pytorch/xla/pull/2117\n        # From pytorch/xla if `torch_xla.distributed.parallel_loader.MpDeviceLoader` is not available\n        def __init__(self, loader, device, **kwargs):\n            self._loader = loader\n            self._device = device\n            self._parallel_loader_kwargs = kwargs\n\n        def __iter__(self):\n            parallel_loader = ParallelLoader(self._loader, [self._device], **self._parallel_loader_kwargs)\n            return parallel_loader.per_device_loader(self._device)\n\n        def __len__(self):\n            return len(self._loader)\n\n    class _XLADistributedOptimizer(Optimizer):\n        def __init__(self, optimizer):\n            super(self.__class__, self).__init__(optimizer.param_groups)\n            self.wrapped_optimizer = optimizer\n\n        def step(self, closure=None):\n            xm.optimizer_step(self.wrapped_optimizer, barrier=True)\n", "path": "ignite/distributed/auto.py"}]}