Gathers variables from different GPUs on a specified device (-1 means the CPU), with dictionary support.
(outputs, target_device, dim=0)
| 25 | |
| 26 | |
| 27 | def dict_gather(outputs, target_device, dim=0): |
| 28 | """ |
| 29 | Gathers variables from different GPUs on a specified device |
| 30 | (-1 means the CPU), with dictionary support. |
| 31 | """ |
| 32 | def gather_map(outputs): |
| 33 | out = outputs[0] |
| 34 | if torch.is_tensor(out): |
| 35 | # MJY(20180330) HACK:: force nr_dims > 0 |
| 36 | if out.dim() == 0: |
| 37 | outputs = [o.unsqueeze(0) for o in outputs] |
| 38 | return Gather.apply(target_device, dim, *outputs) |
| 39 | elif out is None: |
| 40 | return None |
| 41 | elif isinstance(out, collections.Mapping): |
| 42 | return {k: gather_map([o[k] for o in outputs]) for k in out} |
| 43 | elif isinstance(out, collections.Sequence): |
| 44 | return type(out)(map(gather_map, zip(*outputs))) |
| 45 | return gather_map(outputs) |
| 46 | |
| 47 | |
| 48 | class DictGatherDataParallel(nn.DataParallel): |