Fix result gathering with varying tensor shapes

CHANGELOG.md

@@ -146,6 +146,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed an issue that caused `Trainer.test()` to stall in ddp mode ([#2997](https://github.com/PyTorchLightning/pytorch-lightning/pull/2997))
 
+- Fixed gathering of results with tensors of varying shape ([#3020](https://github.com/PyTorchLightning/pytorch-lightning/pull/3020))
+
 ## [0.8.5] - 2020-07-09
 
 ### Added

pytorch_lightning/core/step_result.py

@@ -1,6 +1,6 @@
 import numbers
 from copy import copy
-from typing import Optional, Dict, Union, Sequence, Callable, MutableMapping, Any
+from typing import Optional, Dict, Union, Sequence, Callable, MutableMapping, Any, List, Tuple
 
 import torch
 from torch import Tensor
@@ -405,19 +405,23 @@ def recursive_stack(result: MutableMapping):
         if isinstance(v, dict):
             recursive_stack(v)
 
-        if isinstance(v, list) and len(v) > 0 and isinstance(v[0], torch.Tensor):
-            v = torch.stack(v)
-            result[k] = v
+        result[k] = collate_tensors(v)
 
 
-def recursive_padded_stack(result: MutableMapping):
-    for k, v in result.items():
-        if isinstance(v, dict):
-            recursive_stack(v)
+def collate_tensors(items: Union[List, Tuple]) -> Union[Tensor, List, Tuple]:
+    if not items or not isinstance(items, (list, tuple)) or any(not isinstance(item, Tensor) for item in items):
+        # items is not a sequence, empty, or contains non-tensors
+        return items
+
+    if all(item.ndim == 0 for item in items):
+        # all tensors are scalars, we need to stack
+        return torch.stack(items)
+
+    if all(item.ndim >= 1 and item.shape[1:] == items[0].shape[1:] for item in items):
+        # we can concatenate along the first dimension
+        return torch.cat(items)
 
-        if isinstance(v, list) and len(v) > 0 and isinstance(v[0], torch.Tensor):
-            v = torch.stack(v)
-            result[k] = v
+    return items
 
 
 class TrainResult(Result):

tests/core/test_results.py

@@ -174,3 +174,52 @@ def test_result_obj_predictions_ddp_spawn(tmpdir):
         predictions = torch.load(prediction_file)
         size += len(predictions)
     assert size == len(dm.mnist_test)
+
+
+def test_result_gather_stack():
+    """ Test that tensors get concatenated when they all have the same shape. """
+    outputs = [
+        {"foo": torch.zeros(4, 5)},
+        {"foo": torch.zeros(4, 5)},
+        {"foo": torch.zeros(4, 5)},
+    ]
+    result = Result.gather(outputs)
+    assert isinstance(result["foo"], torch.Tensor)
+    assert list(result["foo"].shape) == [12, 5]
+
+
+def test_result_gather_concatenate():
+    """ Test that tensors get concatenated when they have varying size in first dimension. """
+    outputs = [
+        {"foo": torch.zeros(4, 5)},
+        {"foo": torch.zeros(8, 5)},
+        {"foo": torch.zeros(3, 5)},
+    ]
+    result = Result.gather(outputs)
+    assert isinstance(result["foo"], torch.Tensor)
+    assert list(result["foo"].shape) == [15, 5]
+
+
+def test_result_gather_scalar():
+    """ Test that 0-dim tensors get gathered and stacked correctly. """
+    outputs = [
+        {"foo": torch.tensor(1)},
+        {"foo": torch.tensor(2)},
+        {"foo": torch.tensor(3)},
+    ]
+    result = Result.gather(outputs)
+    assert isinstance(result["foo"], torch.Tensor)
+    assert list(result["foo"].shape) == [3]
+
+
+def test_result_gather_different_shapes():
+    """ Test that tensors of varying shape get gathered into a list. """
+    outputs = [
+        {"foo": torch.tensor(1)},
+        {"foo": torch.zeros(2, 3)},
+        {"foo": torch.zeros(1, 2, 3)},
+    ]
+    result = Result.gather(outputs)
+    expected = [torch.tensor(1), torch.zeros(2, 3), torch.zeros(1, 2, 3)]
+    assert isinstance(result["foo"], list)
+    assert all(torch.eq(r, e).all() for r, e in zip(result["foo"], expected))