data_parallel.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235
  1. import operator
  2. import torch
  3. import warnings
  4. from itertools import chain
  5. from ..modules import Module
  6. from .scatter_gather import scatter_kwargs, gather
  7. from .replicate import replicate
  8. from .parallel_apply import parallel_apply
  9. from torch._utils import (
  10. _get_all_device_indices,
  11. _get_available_device_type,
  12. _get_device_index,
  13. _get_devices_properties
  14. )
  15. __all__ = ['DataParallel', 'data_parallel']
  16. def _check_balance(device_ids):
  17. imbalance_warn = """
  18. There is an imbalance between your GPUs. You may want to exclude GPU {} which
  19. has less than 75% of the memory or cores of GPU {}. You can do so by setting
  20. the device_ids argument to DataParallel, or by setting the CUDA_VISIBLE_DEVICES
  21. environment variable."""
  22. device_ids = [_get_device_index(x, True) for x in device_ids]
  23. dev_props = _get_devices_properties(device_ids)
  24. def warn_imbalance(get_prop):
  25. values = [get_prop(props) for props in dev_props]
  26. min_pos, min_val = min(enumerate(values), key=operator.itemgetter(1))
  27. max_pos, max_val = max(enumerate(values), key=operator.itemgetter(1))
  28. if min_val / max_val < 0.75:
  29. warnings.warn(imbalance_warn.format(device_ids[min_pos], device_ids[max_pos]))
  30. return True
  31. return False
  32. if warn_imbalance(lambda props: props.total_memory):
  33. return
  34. if warn_imbalance(lambda props: props.multi_processor_count):
  35. return
  36. class DataParallel(Module):
  37. r"""Implements data parallelism at the module level.
  38. This container parallelizes the application of the given :attr:`module` by
  39. splitting the input across the specified devices by chunking in the batch
  40. dimension (other objects will be copied once per device). In the forward
  41. pass, the module is replicated on each device, and each replica handles a
  42. portion of the input. During the backwards pass, gradients from each replica
  43. are summed into the original module.
  44. The batch size should be larger than the number of GPUs used.
  45. .. warning::
  46. It is recommended to use :class:`~torch.nn.parallel.DistributedDataParallel`,
  47. instead of this class, to do multi-GPU training, even if there is only a single
  48. node. See: :ref:`cuda-nn-ddp-instead` and :ref:`ddp`.
  49. Arbitrary positional and keyword inputs are allowed to be passed into
  50. DataParallel but some types are specially handled. tensors will be
  51. **scattered** on dim specified (default 0). tuple, list and dict types will
  52. be shallow copied. The other types will be shared among different threads
  53. and can be corrupted if written to in the model's forward pass.
  54. The parallelized :attr:`module` must have its parameters and buffers on
  55. ``device_ids[0]`` before running this :class:`~torch.nn.DataParallel`
  56. module.
  57. .. warning::
  58. In each forward, :attr:`module` is **replicated** on each device, so any
  59. updates to the running module in ``forward`` will be lost. For example,
  60. if :attr:`module` has a counter attribute that is incremented in each
  61. ``forward``, it will always stay at the initial value because the update
  62. is done on the replicas which are destroyed after ``forward``. However,
  63. :class:`~torch.nn.DataParallel` guarantees that the replica on
  64. ``device[0]`` will have its parameters and buffers sharing storage with
  65. the base parallelized :attr:`module`. So **in-place** updates to the
  66. parameters or buffers on ``device[0]`` will be recorded. E.g.,
  67. :class:`~torch.nn.BatchNorm2d` and :func:`~torch.nn.utils.spectral_norm`
  68. rely on this behavior to update the buffers.
  69. .. warning::
  70. Forward and backward hooks defined on :attr:`module` and its submodules
  71. will be invoked ``len(device_ids)`` times, each with inputs located on
  72. a particular device. Particularly, the hooks are only guaranteed to be
  73. executed in correct order with respect to operations on corresponding
  74. devices. For example, it is not guaranteed that hooks set via
  75. :meth:`~torch.nn.Module.register_forward_pre_hook` be executed before
  76. `all` ``len(device_ids)`` :meth:`~torch.nn.Module.forward` calls, but
  77. that each such hook be executed before the corresponding
  78. :meth:`~torch.nn.Module.forward` call of that device.
  79. .. warning::
  80. When :attr:`module` returns a scalar (i.e., 0-dimensional tensor) in
  81. :func:`forward`, this wrapper will return a vector of length equal to
  82. number of devices used in data parallelism, containing the result from
  83. each device.
  84. .. note::
  85. There is a subtlety in using the
  86. ``pack sequence -> recurrent network -> unpack sequence`` pattern in a
  87. :class:`~torch.nn.Module` wrapped in :class:`~torch.nn.DataParallel`.
  88. See :ref:`pack-rnn-unpack-with-data-parallelism` section in FAQ for
  89. details.
  90. Args:
  91. module (Module): module to be parallelized
  92. device_ids (list of int or torch.device): CUDA devices (default: all devices)
  93. output_device (int or torch.device): device location of output (default: device_ids[0])
  94. Attributes:
  95. module (Module): the module to be parallelized
  96. Example::
  97. >>> # xdoctest: +SKIP
  98. >>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
  99. >>> output = net(input_var) # input_var can be on any device, including CPU
  100. """
  101. # TODO: update notes/cuda.rst when this class handles 8+ GPUs well
  102. def __init__(self, module, device_ids=None, output_device=None, dim=0):
  103. super().__init__()
  104. torch._C._log_api_usage_once("torch.nn.parallel.DataParallel")
  105. device_type = _get_available_device_type()
  106. if device_type is None:
  107. self.module = module
  108. self.device_ids = []
  109. return
  110. if device_ids is None:
  111. device_ids = _get_all_device_indices()
  112. if output_device is None:
  113. output_device = device_ids[0]
  114. self.dim = dim
  115. self.module = module
  116. self.device_ids = [_get_device_index(x, True) for x in device_ids]
  117. self.output_device = _get_device_index(output_device, True)
  118. self.src_device_obj = torch.device(device_type, self.device_ids[0])
  119. _check_balance(self.device_ids)
  120. if len(self.device_ids) == 1:
  121. self.module.to(self.src_device_obj)
  122. def forward(self, *inputs, **kwargs):
  123. with torch.autograd.profiler.record_function("DataParallel.forward"):
  124. if not self.device_ids:
  125. return self.module(*inputs, **kwargs)
  126. for t in chain(self.module.parameters(), self.module.buffers()):
  127. if t.device != self.src_device_obj:
  128. raise RuntimeError("module must have its parameters and buffers "
  129. "on device {} (device_ids[0]) but found one of "
  130. "them on device: {}".format(self.src_device_obj, t.device))
  131. inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
  132. # for forward function without any inputs, empty list and dict will be created
  133. # so the module can be executed on one device which is the first one in device_ids
  134. if not inputs and not kwargs:
  135. inputs = ((),)
  136. kwargs = ({},)
  137. if len(self.device_ids) == 1:
  138. return self.module(*inputs[0], **kwargs[0])
  139. replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
  140. outputs = self.parallel_apply(replicas, inputs, kwargs)
  141. return self.gather(outputs, self.output_device)
  142. def replicate(self, module, device_ids):
  143. return replicate(module, device_ids, not torch.is_grad_enabled())
  144. def scatter(self, inputs, kwargs, device_ids):
  145. return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
  146. def parallel_apply(self, replicas, inputs, kwargs):
  147. return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  148. def gather(self, outputs, output_device):
  149. return gather(outputs, output_device, dim=self.dim)
  150. def data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None):
  151. r"""Evaluates module(input) in parallel across the GPUs given in device_ids.
  152. This is the functional version of the DataParallel module.
  153. Args:
  154. module (Module): the module to evaluate in parallel
  155. inputs (Tensor): inputs to the module
  156. device_ids (list of int or torch.device): GPU ids on which to replicate module
  157. output_device (list of int or torch.device): GPU location of the output Use -1 to indicate the CPU.
  158. (default: device_ids[0])
  159. Returns:
  160. a Tensor containing the result of module(input) located on
  161. output_device
  162. """
  163. if not isinstance(inputs, tuple):
  164. inputs = (inputs,) if inputs is not None else ()
  165. device_type = _get_available_device_type()
  166. if device_ids is None:
  167. device_ids = _get_all_device_indices()
  168. if output_device is None:
  169. output_device = device_ids[0]
  170. device_ids = [_get_device_index(x, True) for x in device_ids]
  171. output_device = _get_device_index(output_device, True)
  172. src_device_obj = torch.device(device_type, device_ids[0])
  173. for t in chain(module.parameters(), module.buffers()):
  174. if t.device != src_device_obj:
  175. raise RuntimeError("module must have its parameters and buffers "
  176. "on device {} (device_ids[0]) but found one of "
  177. "them on device: {}".format(src_device_obj, t.device))
  178. inputs, module_kwargs = scatter_kwargs(inputs, module_kwargs, device_ids, dim)
  179. # for module without any inputs, empty list and dict will be created
  180. # so the module can be executed on one device which is the first one in device_ids
  181. if not inputs and not module_kwargs:
  182. inputs = ((),)
  183. module_kwargs = ({},)
  184. if len(device_ids) == 1:
  185. return module(*inputs[0], **module_kwargs[0])
  186. used_device_ids = device_ids[:len(inputs)]
  187. replicas = replicate(module, used_device_ids)
  188. outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
  189. return gather(outputs, output_device, dim)