Repository: pytorch/pytorch

Issue:
DataParallel Gather works only with iterable outputs
I'm running a network across multiple GPUs and pass the input data through a dictionary. The Variables stored as items of the input dictionary are properly scattered across the batch dimension, and the forward pass terminates correctly.

However, when returning the output as a dictionary I get the following runtime error:
```
torch/nn/parallel/scatter_gather.pyc in gather_map(outputs)
     47         if out is None:
     48             return None
---> 49         return type(out)(map(gather_map, zip(*outputs)))
     50     return gather_map(outputs)
torch/nn/parallel/scatter_gather.pyc in gather_map(outputs)
     43     def gather_map(outputs):
     44         out = outputs[0]
---> 45         if isinstance(out, Variable):
     46             return Gather(target_device, dim=dim)(*outputs)
     47         if out is None:
RuntimeError: maximum recursion depth exceeded in __instancecheck__
```

The reason is that the function `gather_map` in [scatter_gather.py](https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/scatter_gather.py#L46-L52) only supports Variables or iterables of variables as its input.

However, `scatter_map` in [scatter_gather.py](https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/scatter_gather.py#L21-L23) also supports dictionaries.

Is there a reason for this discrepancy? Would it be useful if I made a pull request and added this functionality?

I am implementing a network with multiple sub-networks whose outputs may or may not be computed (based on a config file) and it would be useful to be able to pass all of them in a compact way out of the data parallel wrapper.

**SNIPPET REPLICATING ERROR:**

```
import torch
import torch.nn as nn
from torch.autograd import Variable

class MyModel(nn.Module):

    def __init__(self):
        super(MyModel, self).__init__()
        self.block1 = nn.Linear(10, 20)
        self.block2 = nn.Linear(20, 20)

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        return x

class MyModelDictInput(nn.Module):

    def __init__(self):
        super(MyModelDictInput, self).__init__()
        self.block1 = nn.Linear(10, 20)
        self.block2 = nn.Linear(20, 20)

    def forward(self, d):
        x = d['an_input']
        x = self.block1(x)
        x = self.block2(x)
        return x

class MyModelDictOutput(nn.Module):

    def __init__(self):
        super(MyModelDictOutput, self).__init__()
        self.block1 = nn.Linear(10, 20)
        self.block2 = nn.Linear(20, 20)

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        d = dict()
        d['an_output'] = x
        return d

# create random input
i = Variable(torch.rand((4,10)))
d = {'an_input':i}

# example 1:
print('input is a Variable, output is a Variable')
net = nn.DataParallel(MyModel()).cuda()
o   = net.forward(i)
print(o)

# example 2:
print('input is a dict, output is a Variable')
net = nn.DataParallel(MyModelDictInput()).cuda()
o   = net.forward(d)
print(o)

# example 3:
print('input is a Variable, output is a dict')
net = nn.DataParallel(MyModelDictOutput()).cuda()
o   = net.forward(i)
print(o)
```


Discussion:
Has anyone ever taken a look at this? Thanks!
This would help simplifying our framework.
We are now using the method mentioned in [Training network with multiple outputs with multi gpus](https://discuss.pytorch.org/t/training-network-with-multiple-outputs-with-multi-gpus/6344), but it is a bit hacky.

Resolving patch:
diff --git a/torch/nn/parallel/scatter_gather.py b/torch/nn/parallel/scatter_gather.py
--- a/torch/nn/parallel/scatter_gather.py
+++ b/torch/nn/parallel/scatter_gather.py
@@ -57,6 +57,11 @@ def gather_map(outputs):
             return Gather.apply(target_device, dim, *outputs)
         if out is None:
             return None
+        if isinstance(out, dict):
+            if not all((len(out) == len(d) for d in outputs)):
+                raise ValueError('All dicts must have the same number of keys')
+            return type(out)(((k, gather_map([d[k] for d in outputs]))
+                              for k in out))
         return type(out)(map(gather_map, zip(*outputs)))
 
     # Recursive function calls like this create reference cycles.


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -2054,6 +2054,18 @@ def test_gather_cpu(self):
     def test_gather_gpu(self):
         self._test_gather(0)
 
+    @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
+    def test_gather_different_len_dicts(self):
+        inputs = (
+            {'a': Variable(torch.randn(1, 2).cuda(0), requires_grad=True)},
+            {
+                'b': Variable(torch.randn(1, 2).cuda(1), requires_grad=True),
+                'a': Variable(torch.randn(1, 2).cuda(1), requires_grad=True)
+            }
+        )
+        with self.assertRaises(ValueError):
+            _ = dp.gather(inputs, target_device=0)
+
     def _test_broadcast_double_backwards(self, *tensors):
         variables = tuple(Variable(t, requires_grad=True) for t in tensors)
         _assertGradAndGradgradChecks(self, lambda *i: Broadcast.apply((0, 1), *i), variables)
@@ -2297,7 +2309,10 @@ def test_data_parallel_sparse(self):
     @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
     def test_data_parallel_nested_output(self):
         def fn(input):
-            return [input, (input.sin(), input.cos(), [input.add(1)]), input]
+            return [
+                input, (input.sin(), input.cos(), [input.add(1)]), input,
+                {'a': input, 'b': [input.sin()]}
+            ]
 
         class Net(nn.Module):
             def forward(self, input):
@@ -2314,6 +2329,13 @@ def forward(self, input):
         self.assertIsInstance(output[1][2], list)
         self.assertIsInstance(output[1][2][0], Variable)
         self.assertIsInstance(output[2], Variable)
+        self.assertIsInstance(output[3], dict)
+        self.assertEqual(len(output[3]), 2)
+        self.assertIn('a', output[3])
+        self.assertIn('b', output[3])
+        self.assertIsInstance(output[3]['a'], Variable)
+        self.assertIsInstance(output[3]['b'], list)
+        self.assertIsInstance(output[3]['b'][0], Variable)
 
     @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
     def test_data_parallel_nested_input(self):