r"""Apply each `module` in :attr:`modules` in parallel on each of :attr:`devices`. Args: modules (Module): modules to be parallelized inputs (tensor): inputs to the modules devices (list of int or torch.device): CUDA devices :attr:`modules`, :attr:`inputs`, :attr:`k
(
modules: Sequence[Module],
inputs: Sequence[Any],
kwargs_tup: Optional[Sequence[Dict[str, Any]]] = None,
devices: Optional[Sequence[Optional[Union[int, torch.device]]]] = None,
)
| 23 | return None |
| 24 | |
| 25 | def parallel_apply( |
| 26 | modules: Sequence[Module], |
| 27 | inputs: Sequence[Any], |
| 28 | kwargs_tup: Optional[Sequence[Dict[str, Any]]] = None, |
| 29 | devices: Optional[Sequence[Optional[Union[int, torch.device]]]] = None, |
| 30 | ) -> List[Any]: |
| 31 | r"""Apply each `module` in :attr:`modules` in parallel on each of :attr:`devices`. |
| 32 | |
| 33 | Args: |
| 34 | modules (Module): modules to be parallelized |
| 35 | inputs (tensor): inputs to the modules |
| 36 | devices (list of int or torch.device): CUDA devices |
| 37 | |
| 38 | :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and |
| 39 | :attr:`devices` (if given) should all have same length. Moreover, each |
| 40 | element of :attr:`inputs` can either be a single object as the only argument |
| 41 | to a module, or a collection of positional arguments. |
| 42 | """ |
| 43 | assert len(modules) == len(inputs), f'The number of modules {len(modules)} is not equal to the number of inputs {len(inputs)}' |
| 44 | if kwargs_tup is not None: |
| 45 | assert len(modules) == len(kwargs_tup) |
| 46 | else: |
| 47 | kwargs_tup = (cast(Dict[str, Any], {}),) * len(modules) |
| 48 | if devices is not None: |
| 49 | assert len(modules) == len(devices) |
| 50 | else: |
| 51 | devices = [None] * len(modules) |
| 52 | devices = [_get_device_index(x, True) for x in devices] |
| 53 | streams = [torch.cuda.current_stream(x) for x in devices] |
| 54 | lock = threading.Lock() |
| 55 | results = {} |
| 56 | grad_enabled, autocast_enabled = torch.is_grad_enabled(), torch.is_autocast_enabled() |
| 57 | |
| 58 | def _worker( |
| 59 | i: int, |
| 60 | module: Module, |
| 61 | input: Any, |
| 62 | kwargs: Dict[str, Any], |
| 63 | device: Optional[Union[int, torch.device]] = None, |
| 64 | stream: Optional[torch.cuda.Stream] = None, |
| 65 | ) -> None: |
| 66 | torch.set_grad_enabled(grad_enabled) |
| 67 | if device is None: |
| 68 | t = get_a_var(input) |
| 69 | if t is None: |
| 70 | with lock: |
| 71 | results[i] = ExceptionWrapper( |
| 72 | where=f"in replica {i}, no device was provided and no tensor input was found; " |
| 73 | "device cannot be resolved") |
| 74 | return |
| 75 | device = t.get_device() |
| 76 | if stream is None: |
| 77 | stream = torch.cuda.current_stream(device) |
| 78 | try: |
| 79 | with torch.cuda.device(device), torch.cuda.stream(stream), autocast(enabled=autocast_enabled): |
| 80 | # this also avoids accidental slicing of `input` if it is a Tensor |
| 81 | if not isinstance(input, (list, tuple)): |
| 82 | input = (input,) |
no test coverage detected
searching dependent graphs…