| | |
| |
|
| | import torch.cuda as cuda |
| | import torch.nn as nn |
| | import torch |
| | import collections |
| | from torch.nn.parallel._functions import Gather |
| |
|
| |
|
| | __all__ = ['UserScatteredDataParallel', 'user_scattered_collate', 'async_copy_to'] |
| |
|
| |
|
| | def async_copy_to(obj, dev, main_stream=None): |
| | if torch.is_tensor(obj): |
| | v = obj.cuda(dev, non_blocking=True) |
| | if main_stream is not None: |
| | v.data.record_stream(main_stream) |
| | return v |
| | elif isinstance(obj, collections.Mapping): |
| | return {k: async_copy_to(o, dev, main_stream) for k, o in obj.items()} |
| | elif isinstance(obj, collections.Sequence): |
| | return [async_copy_to(o, dev, main_stream) for o in obj] |
| | else: |
| | return obj |
| |
|
| |
|
| | def dict_gather(outputs, target_device, dim=0): |
| | """ |
| | Gathers variables from different GPUs on a specified device |
| | (-1 means the CPU), with dictionary support. |
| | """ |
| | def gather_map(outputs): |
| | out = outputs[0] |
| | if torch.is_tensor(out): |
| | |
| | if out.dim() == 0: |
| | outputs = [o.unsqueeze(0) for o in outputs] |
| | return Gather.apply(target_device, dim, *outputs) |
| | elif out is None: |
| | return None |
| | elif isinstance(out, collections.Mapping): |
| | return {k: gather_map([o[k] for o in outputs]) for k in out} |
| | elif isinstance(out, collections.Sequence): |
| | return type(out)(map(gather_map, zip(*outputs))) |
| | return gather_map(outputs) |
| |
|
| |
|
| | class DictGatherDataParallel(nn.DataParallel): |
| | def gather(self, outputs, output_device): |
| | return dict_gather(outputs, output_device, dim=self.dim) |
| |
|
| |
|
| | class UserScatteredDataParallel(DictGatherDataParallel): |
| | def scatter(self, inputs, kwargs, device_ids): |
| | assert len(inputs) == 1 |
| | inputs = inputs[0] |
| | inputs = _async_copy_stream(inputs, device_ids) |
| | inputs = [[i] for i in inputs] |
| | assert len(kwargs) == 0 |
| | kwargs = [{} for _ in range(len(inputs))] |
| |
|
| | return inputs, kwargs |
| |
|
| |
|
| | def user_scattered_collate(batch): |
| | return batch |
| |
|
| |
|
| | def _async_copy(inputs, device_ids): |
| | nr_devs = len(device_ids) |
| | assert type(inputs) in (tuple, list) |
| | assert len(inputs) == nr_devs |
| |
|
| | outputs = [] |
| | for i, dev in zip(inputs, device_ids): |
| | with cuda.device(dev): |
| | outputs.append(async_copy_to(i, dev)) |
| |
|
| | return tuple(outputs) |
| |
|
| |
|
| | def _async_copy_stream(inputs, device_ids): |
| | nr_devs = len(device_ids) |
| | assert type(inputs) in (tuple, list) |
| | assert len(inputs) == nr_devs |
| |
|
| | outputs = [] |
| | streams = [_get_stream(d) for d in device_ids] |
| | for i, dev, stream in zip(inputs, device_ids, streams): |
| | with cuda.device(dev): |
| | main_stream = cuda.current_stream() |
| | with cuda.stream(stream): |
| | outputs.append(async_copy_to(i, dev, main_stream=main_stream)) |
| | main_stream.wait_stream(stream) |
| |
|
| | return outputs |
| |
|
| |
|
| | """Adapted from: torch/nn/parallel/_functions.py""" |
| | |
| | _streams = None |
| |
|
| |
|
| | def _get_stream(device): |
| | """Gets a background stream for copying between CPU and GPU""" |
| | global _streams |
| | if device == -1: |
| | return None |
| | if _streams is None: |
| | _streams = [None] * cuda.device_count() |
| | if _streams[device] is None: _streams[device] = cuda.Stream(device) |
| | return _streams[device] |
| |
|
