Evaluates module(input) in parallel across the GPUs given in device_ids. This is the functional version of the DataParallel module. Args: module: the module to evaluate in parallel inputs: inputs to the module device_ids: GPU ids on which to replicate module
(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None, gather=True)
| 69 | |
| 70 | |
| 71 | def data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None, gather=True): |
| 72 | """ |
| 73 | Evaluates module(input) in parallel across the GPUs given in device_ids. |
| 74 | This is the functional version of the DataParallel module. |
| 75 | Args: |
| 76 | module: the module to evaluate in parallel |
| 77 | inputs: inputs to the module |
| 78 | device_ids: GPU ids on which to replicate module |
| 79 | output_device: GPU location of the output Use -1 to indicate the CPU. |
| 80 | (default: device_ids[0]) |
| 81 | Returns: |
| 82 | a Tensor containing the result of module(input) located on |
| 83 | output_device |
| 84 | """ |
| 85 | if not isinstance(inputs, tuple): |
| 86 | inputs = (inputs,) |
| 87 | |
| 88 | if device_ids is None: |
| 89 | device_ids = list(range(torch.cuda.device_count())) |
| 90 | |
| 91 | if output_device is None: |
| 92 | output_device = device_ids[0] |
| 93 | |
| 94 | inputs, module_kwargs = scatter_kwargs(inputs, module_kwargs, device_ids, dim) |
| 95 | if len(device_ids) == 1: |
| 96 | return module(*inputs[0], **module_kwargs[0]) |
| 97 | used_device_ids = device_ids[:len(inputs)] |
| 98 | replicas = replicate(module, used_device_ids) |
| 99 | outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids) |
| 100 | if gather: |
| 101 | return gather(outputs, output_device, dim) |
| 102 | else: |
| 103 | return outputs |
| 104 | |
| 105 | |
| 106 |
nothing calls this directly
no outgoing calls
no test coverage detected