Repository: Project-MONAI/MONAI

Issue:
Should detach the Tensor before decollate_batch
**Describe the bug**
Currently, we decollate a Tensor data(like: pred) from batch to list without `detach()` first.
I will enhance it soon.




Discussion:

Patch:
diff --git a/monai/data/utils.py b/monai/data/utils.py
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -15,10 +15,10 @@
 import os
 import pickle
 import warnings
-from collections import abc, defaultdict
+from collections import defaultdict
 from itertools import product, starmap
 from pathlib import PurePath
-from typing import Any, Dict, Generator, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import Dict, Generator, Iterable, List, Mapping, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
@@ -36,7 +36,6 @@
     optional_import,
 )
 from monai.utils.enums import Method
-from monai.utils.misc import issequenceiterable
 
 nib, _ = optional_import("nibabel")
 
@@ -254,10 +253,10 @@ def list_data_collate(batch: Sequence):
     """
     elem = batch[0]
     data = [i for k in batch for i in k] if isinstance(elem, list) else batch
+    key = None
     try:
         elem = batch[0]
-        key = None
-        if isinstance(elem, abc.Mapping):
+        if isinstance(elem, Mapping):
             ret = {}
             for k in elem:
                 key = k
@@ -288,17 +287,18 @@ def list_data_collate(batch: Sequence):
         raise TypeError(re_str)
 
 
-def decollate_batch(data: Union[dict, list, torch.Tensor], batch_size: Optional[int] = None) -> List[dict]:
+def decollate_batch(batch, detach=True):
     """De-collate a batch of data (for example, as produced by a `DataLoader`).
 
-    Returns a list of dictionaries, list or Tensor, mapping to a given batch.
+    Returns a list of structures with the original tensor's 0-th dimension sliced into elements using `torch.unbind`.
 
     Images originally stored as (B,C,H,W,[D]) will be returned as (C,H,W,[D]). Other information,
     such as metadata, may have been stored in a list (or a list inside nested dictionaries). In
     this case we return the element of the list corresponding to the batch idx.
 
     Return types aren't guaranteed to be the same as the original, since numpy arrays will have been
-    converted to torch.Tensor, and tuples/lists may have been converted to lists of tensors
+    converted to torch.Tensor, sequences may be converted to lists of tensors,
+    mappings may be converted into dictionaries.
 
     For example:
 
@@ -315,63 +315,50 @@ def decollate_batch(data: Union[dict, list, torch.Tensor], batch_size: Optional[
         print(out[0])
         >>> {'image': tensor([[[4.3549e-01...43e-01]]]), 'image_meta_dict': {'scl_slope': 0.0}}
 
-        batch_data = [torch.rand((2,1,10,10), torch.rand((2,3,5,5)]
+        batch_data = [torch.rand((2,1,10,10)), torch.rand((2,3,5,5))]
         out = decollate_batch(batch_data)
         print(out[0])
         >>> [tensor([[[4.3549e-01...43e-01]]], tensor([[[5.3435e-01...45e-01]]])]
 
-        batch_data = torch.rand((2,1,10,10)
+        batch_data = torch.rand((2,1,10,10))
         out = decollate_batch(batch_data)
         print(out[0])
         >>> tensor([[[4.3549e-01...43e-01]]])
 
     Args:
-        data: data to be de-collated.
-        batch_size: number of batches in data. If `None` is passed, try to figure out batch size.
-    """
-
-    def torch_to_single(d: torch.Tensor):
-        """If input is a torch.Tensor with only 1 element, return just the element."""
-        return d if d.numel() > 1 else d.item()
-
-    def decollate(data: Any, idx: int):
-        """Recursively de-collate."""
-        if isinstance(data, dict):
-            return {k: decollate(v, idx) for k, v in data.items()}
-        if isinstance(data, torch.Tensor):
-            out = data[idx]
-            return torch_to_single(out)
-        if isinstance(data, list):
-            if len(data) == 0:
-                return data
-            if isinstance(data[0], torch.Tensor):
-                return [torch_to_single(d[idx]) for d in data]
-            if issequenceiterable(data[0]):
-                return [decollate(d, idx) for d in data]
-            return data[idx]
-        raise TypeError(f"Not sure how to de-collate type: {type(data)}")
-
-    def _detect_batch_size(batch_data):
-        for v in batch_data:
-            if isinstance(v, torch.Tensor):
-                return v.shape[0]
-        warnings.warn("batch_data is not a sequence of tensors in decollate, use `len(batch_data[0])` directly.")
-        return len(batch_data[0])
-
-    result: List[Any]
-    if isinstance(data, dict):
-        batch_size = _detect_batch_size(batch_data=data.values()) if batch_size is None else batch_size
-        result = [{key: decollate(data[key], idx) for key in data.keys()} for idx in range(batch_size)]
-    elif isinstance(data, list):
-        batch_size = _detect_batch_size(batch_data=data) if batch_size is None else batch_size
-        result = [[decollate(d, idx) for d in data] for idx in range(batch_size)]
-    elif isinstance(data, torch.Tensor):
-        batch_size = data.shape[0]
-        result = [data[idx] for idx in range(batch_size)]
-    else:
-        raise NotImplementedError("Only currently implemented for dictionary, list or Tensor data.")
-
-    return result
+        batch: data to be de-collated.
+        detach: whether to detach the tensors. Scalars tensors will be detached into number types
+            instead of torch tensors.
+    """
+    if batch is None:
+        return batch
+    if isinstance(batch, (float, int, str, bytes)):
+        return batch
+    if isinstance(batch, torch.Tensor):
+        if detach:
+            batch = batch.detach()
+        if batch.ndim == 0:
+            return batch.item() if detach else batch
+        out_list = torch.unbind(batch, dim=0)
+        if out_list[0].ndim == 0 and detach:
+            return [t.item() for t in out_list]
+        return list(out_list)
+    if isinstance(batch, Mapping):
+        _dict_list = {key: decollate_batch(batch[key], detach) for key in batch}
+        return [dict(zip(_dict_list, item)) for item in zip(*_dict_list.values())]
+    if isinstance(batch, Iterable):
+        item_0 = first(batch)
+        if (
+            not isinstance(item_0, Iterable)
+            or isinstance(item_0, (str, bytes))
+            or (isinstance(item_0, torch.Tensor) and item_0.ndim == 0)
+        ):
+            # Not running the usual list decollate here:
+            # don't decollate ['test', 'test'] into [['t', 't'], ['e', 'e'], ['s', 's'], ['t', 't']]
+            # torch.tensor(0) is iterable but iter(torch.tensor(0)) raises TypeError: iteration over a 0-d tensor
+            return [decollate_batch(b, detach) for b in batch]
+        return [list(item) for item in zip(*(decollate_batch(b, detach) for b in batch))]
+    raise NotImplementedError(f"Unable to de-collate: {batch}, type: {type(batch)}.")
 
 
 def pad_list_data_collate(
diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py
--- a/monai/transforms/post/dictionary.py
+++ b/monai/transforms/post/dictionary.py
@@ -351,12 +351,11 @@ class Decollated(MapTransform):
             it fails to determine it automatically.
     """
 
-    def __init__(self, batch_size: Optional[int] = None) -> None:
-        super().__init__(None)
-        self.batch_size = batch_size
+    def __init__(self, keys="") -> None:
+        super().__init__(keys=keys)
 
-    def __call__(self, data: dict) -> List[dict]:
-        return decollate_batch(data, self.batch_size)
+    def __call__(self, data: dict):
+        return decollate_batch(data)
 
 
 class ProbNMSd(MapTransform):
diff --git a/monai/utils/misc.py b/monai/utils/misc.py
--- a/monai/utils/misc.py
+++ b/monai/utils/misc.py
@@ -83,7 +83,7 @@ def issequenceiterable(obj: Any) -> bool:
     """
     if isinstance(obj, torch.Tensor):
         return int(obj.dim()) > 0  # a 0-d tensor is not iterable
-    return isinstance(obj, collections.abc.Iterable) and not isinstance(obj, str)
+    return isinstance(obj, collections.abc.Iterable) and not isinstance(obj, (str, bytes))
 
 
 def ensure_tuple(vals: Any) -> Tuple[Any, ...]:


Test patch:
diff --git a/tests/test_decollate.py b/tests/test_decollate.py
--- a/tests/test_decollate.py
+++ b/tests/test_decollate.py
@@ -56,6 +56,28 @@
 TESTS_LIST.append((RandAffine(prob=0.0, translate_range=10),))
 
 
+TEST_BASIC = [
+    [("channel", "channel"), ["channel", "channel"]],
+    [torch.Tensor([1, 2, 3]), [torch.tensor(1.0), torch.tensor(2.0), torch.tensor(3.0)]],
+    [
+        [[torch.Tensor((1.0, 2.0, 3.0)), torch.Tensor((2.0, 3.0, 1.0))]],
+        [
+            [[torch.tensor(1.0), torch.tensor(2.0)]],
+            [[torch.tensor(2.0), torch.tensor(3.0)]],
+            [[torch.tensor(3.0), torch.tensor(1.0)]],
+        ],
+    ],
+    [torch.Tensor((True, True, False, False)), [1.0, 1.0, 0.0, 0.0]],
+    [
+        [torch.Tensor([1]), torch.Tensor([2]), torch.Tensor([3])],
+        [[torch.tensor(1.0), torch.tensor(2.0), torch.tensor(3.0)]],
+    ],
+    [[None, None], [None, None]],
+    [["test"], ["test"]],
+    [[], []],
+]
+
+
 class _ListCompose(Compose):
     def __call__(self, input_):
         img, metadata = self.transforms[0](input_)
@@ -141,5 +163,53 @@ def test_decollation_list(self, *transforms):
         self.check_decollate(dataset=dataset)
 
 
+class TestBasicDeCollate(unittest.TestCase):
+    @parameterized.expand(TEST_BASIC)
+    def test_decollation_examples(self, input_val, expected_out):
+        out = decollate_batch(input_val)
+        self.assertListEqual(expected_out, out)
+
+    def test_dict_examples(self):
+        test_case = {
+            "meta": {"out": ["test", "test"]},
+            "image_meta_dict": {"scl_slope": torch.Tensor((0.0, 0.0))},
+        }
+        out = decollate_batch(test_case)
+        self.assertEqual(out[0]["meta"]["out"], "test")
+        self.assertEqual(out[0]["image_meta_dict"]["scl_slope"], 0.0)
+
+        test_case = [torch.ones((2, 1, 10, 10)), torch.ones((2, 3, 5, 5))]
+        out = decollate_batch(test_case)
+        self.assertTupleEqual(out[0][0].shape, (1, 10, 10))
+        self.assertTupleEqual(out[0][1].shape, (3, 5, 5))
+
+        test_case = torch.rand((2, 1, 10, 10))
+        out = decollate_batch(test_case)
+        self.assertTupleEqual(out[0].shape, (1, 10, 10))
+
+        test_case = [torch.tensor(0), torch.tensor(0)]
+        out = decollate_batch(test_case, detach=True)
+        self.assertListEqual([0, 0], out)
+        self.assertFalse(isinstance(out[0], torch.Tensor))
+
+        test_case = {"a": [torch.tensor(0), torch.tensor(0)]}
+        out = decollate_batch(test_case, detach=False)
+        self.assertListEqual([{"a": torch.tensor(0)}, {"a": torch.tensor(0)}], out)
+        self.assertTrue(isinstance(out[0]["a"], torch.Tensor))
+
+        test_case = [torch.tensor(0), torch.tensor(0)]
+        out = decollate_batch(test_case, detach=False)
+        self.assertListEqual(test_case, out)
+
+        test_case = {
+            "image": torch.tensor([[[1, 2]], [[3, 4]]]),
+            "label": torch.tensor([[[5, 6]], [[7, 8]]]),
+            "pred": torch.tensor([[[9, 10]], [[11, 12]]]),
+            "out": ["test"],
+        }
+        out = decollate_batch(test_case, detach=False)
+        self.assertEqual(out[0]["out"], "test")
+
+
 if __name__ == "__main__":
     unittest.main()