From 1330f2f9b2881a1c733a8e32ba55c7bda02d6f5c Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Fri, 19 Nov 2021 22:12:24 +0000 Subject: [PATCH 1/4] option to zip_longest Signed-off-by: Wenqi Li --- monai/data/utils.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/monai/data/utils.py b/monai/data/utils.py index 29dbeeda24..c9e8f24689 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -18,7 +18,7 @@ from collections import defaultdict from copy import deepcopy from functools import reduce -from itertools import product, starmap +from itertools import product, starmap, zip_longest from pathlib import PurePath from typing import Any, Dict, Generator, Iterable, List, Mapping, Optional, Sequence, Tuple, Union @@ -297,7 +297,7 @@ def list_data_collate(batch: Sequence): raise TypeError(re_str) from re -def decollate_batch(batch, detach: bool = True): +def decollate_batch(batch, detach: bool = True, pad_zip=True, fillvalue=None): """De-collate a batch of data (for example, as produced by a `DataLoader`). Returns a list of structures with the original tensor's 0-th dimension sliced into elements using `torch.unbind`. @@ -339,6 +339,8 @@ def decollate_batch(batch, detach: bool = True): batch: data to be de-collated. detach: whether to detach the tensors. Scalars tensors will be detached into number types instead of torch tensors. + pad_zip: when the items in a batch indicate different batch size, whether to pad up to the longest. + fillvalue: when `pad_zip` is True, the `fillvalue` to use when padding. """ if batch is None: return batch @@ -355,6 +357,8 @@ def decollate_batch(batch, detach: bool = True): return list(out_list) if isinstance(batch, Mapping): _dict_list = {key: decollate_batch(batch[key], detach) for key in batch} + if pad_zip: + return [dict(zip(_dict_list, item)) for item in zip_longest(*_dict_list.values(), fillvalue=fillvalue)] return [dict(zip(_dict_list, item)) for item in zip(*_dict_list.values())] if isinstance(batch, Iterable): item_0 = first(batch) @@ -367,6 +371,10 @@ def decollate_batch(batch, detach: bool = True): # 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] + if pad_zip: + return [ + list(item) for item in zip_longest(*(decollate_batch(b, detach) for b in batch), fillvalue=fillvalue) + ] 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)}.") From b5c1bbbced5e9530878dd233c07f46257409b6e2 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Sat, 20 Nov 2021 13:49:12 +0000 Subject: [PATCH 2/4] adds a pad option Signed-off-by: Wenqi Li --- monai/data/__init__.py | 1 - monai/data/utils.py | 116 ++++++++------------ monai/transforms/inverse_batch_transform.py | 35 ++++-- tests/test_decollate.py | 21 +++- 4 files changed, 93 insertions(+), 80 deletions(-) diff --git a/monai/data/__init__.py b/monai/data/__init__.py index b2dbdeaaa0..e7fa2b3107 100644 --- a/monai/data/__init__.py +++ b/monai/data/__init__.py @@ -62,7 +62,6 @@ partition_dataset_classes, pickle_hashing, rectify_header_sform_qform, - rep_scalar_to_batch, resample_datalist, select_cross_validation_folds, set_rnd, diff --git a/monai/data/utils.py b/monai/data/utils.py index 9e8d31334f..51e908c9ab 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -73,7 +73,6 @@ "pickle_hashing", "sorted_dict", "decollate_batch", - "rep_scalar_to_batch", "pad_list_data_collate", "no_collation", "convert_tables_to_dicts", @@ -297,7 +296,24 @@ def list_data_collate(batch: Sequence): raise TypeError(re_str) from re -def decollate_batch(batch, detach: bool = True, pad_zip=True, fillvalue=None): +def _non_zipping_check(data): + """ + util used by decollate batch, to identify batch size from the collated data. + returns batch_size and the list of non-iterable items. + """ + b, non_iterable = 1, [] + for k, v in data.items() if isinstance(data, Mapping) else enumerate(data): + if not isinstance(v, Iterable) or isinstance(v, (str, bytes)) or (isinstance(v, torch.Tensor) and v.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 + non_iterable.append(k) + elif hasattr(v, "__len__"): + b = max(b, len(v)) + return b, non_iterable + + +def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): """De-collate a batch of data (for example, as produced by a `DataLoader`). Returns a list of structures with the original tensor's 0-th dimension sliced into elements using `torch.unbind`. @@ -335,12 +351,23 @@ def decollate_batch(batch, detach: bool = True, pad_zip=True, fillvalue=None): print(out[0]) >>> tensor([[[4.3549e-01...43e-01]]]) + batch_data = { + "image": [1, 2, 3], "meta": [4, 5], # undetermined batch size + } + out = decollate_batch(batch_data, pad=True, fill_value=0) + print(out) + >>> [{'image': 1, 'meta': 4}, {'image': 2, 'meta': 5}, {'image': 3, 'meta': 0}] + out = decollate_batch(batch_data, pad=False) + print(out) + >>> [{'image': 1, 'meta': 4}, {'image': 2, 'meta': 5}] + Args: batch: data to be de-collated. detach: whether to detach the tensors. Scalars tensors will be detached into number types instead of torch tensors. - pad_zip: when the items in a batch indicate different batch size, whether to pad up to the longest. - fillvalue: when `pad_zip` is True, the `fillvalue` to use when padding. + pad: when the items in a batch indicate different batch size, whether to pad all the sequences to the longest. + If False, the batch size will be the length of the shortest sequence. + fill_value: when `pad` is True, the `fillvalue` to use when padding. """ if batch is None: return batch @@ -356,75 +383,28 @@ def decollate_batch(batch, detach: bool = True, pad_zip=True, fillvalue=None): 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} - if pad_zip: - return [dict(zip(_dict_list, item)) for item in zip_longest(*_dict_list.values(), fillvalue=fillvalue)] + _dict_list = {key: decollate_batch(batch[key], detach, pad=pad, fill_value=fill_value) for key in batch} + if pad: + b, non_iterable = _non_zipping_check(_dict_list) + if len(non_iterable) == len(_dict_list): + return _dict_list + for k in non_iterable: + _dict_list[k] = [deepcopy(_dict_list[k]) for _ in range(b)] + return [dict(zip(_dict_list, item)) for item in zip_longest(*_dict_list.values(), fillvalue=fill_value)] 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] - if pad_zip: - return [ - list(item) for item in zip_longest(*(decollate_batch(b, detach) for b in batch), fillvalue=fillvalue) - ] - return [list(item) for item in zip(*(decollate_batch(b, detach) for b in batch))] + _lists = [decollate_batch(b, detach, pad=pad, fill_value=fill_value) for b in batch] + b, non_iterable = _non_zipping_check(_lists) # type: ignore + if len(non_iterable) == len(_lists): # all non-iterable + return _lists + if pad: + for k in non_iterable: + _lists[k] = [deepcopy(_lists[k]) for _ in range(b)] + return [list(item) for item in zip_longest(*_lists, fillvalue=fill_value)] + return [list(item) for item in zip(*_lists)] raise NotImplementedError(f"Unable to de-collate: {batch}, type: {type(batch)}.") -def rep_scalar_to_batch(batch_data: Union[List, Dict]) -> Union[List, Dict]: - """ - Utility tp replicate the scalar items of a list or dictionary to ensure all the items have batch dimension. - It leverages `decollate_batch(detach=False)` to filter out the scalar items. - - """ - - def _detect_batch_size(batch_data: Sequence): - """ - Detect the batch size from a list of data, some items in the list have batch dim, some not. - - """ - for v in batch_data: - if isinstance(v, torch.Tensor) and v.ndim > 0: - return v.shape[0] - for v in batch_data: - if issequenceiterable(v): - warnings.warn("batch_data doesn't contain batched Tensor data, use the length of first sequence data.") - return len(v) - raise RuntimeError("failed to automatically detect the batch size.") - - if isinstance(batch_data, dict): - batch_size = _detect_batch_size(list(batch_data.values())) - dict_batch = {} - for k, v in batch_data.items(): - if decollate_batch(v, detach=False) == v and not isinstance(v, list): - # if decollating a list, the result may be the same list, so should skip this case - dict_batch[k] = [deepcopy(decollate_batch(v, detach=True)) for _ in range(batch_size)] - else: - dict_batch[k] = v - - return dict_batch - if isinstance(batch_data, list): - batch_size = _detect_batch_size(batch_data) - list_batch = [] - for b in batch_data: - if decollate_batch(b, detach=False) == b and not isinstance(b, list): - list_batch.append([deepcopy(decollate_batch(b, detach=True)) for _ in range(batch_size)]) - else: - list_batch.append(b) - - return list_batch - # if not dict or list, just return the original data - return batch_data - - def pad_list_data_collate( batch: Sequence, method: Union[Method, str] = Method.SYMMETRIC, diff --git a/monai/transforms/inverse_batch_transform.py b/monai/transforms/inverse_batch_transform.py index e220d5f350..9d9d8c2a5d 100644 --- a/monai/transforms/inverse_batch_transform.py +++ b/monai/transforms/inverse_batch_transform.py @@ -17,7 +17,7 @@ from monai.config import KeysCollection from monai.data.dataloader import DataLoader -from monai.data.utils import decollate_batch, no_collation, pad_list_data_collate, rep_scalar_to_batch +from monai.data.utils import decollate_batch, no_collation, pad_list_data_collate from monai.transforms.croppad.batch import PadListDataCollate from monai.transforms.inverse import InvertibleTransform from monai.transforms.transform import MapTransform, Transform @@ -60,6 +60,8 @@ def __init__( collate_fn: Optional[Callable] = no_collation, num_workers: Optional[int] = 0, detach: bool = True, + pad_batch: bool = True, + fill_value=None, ) -> None: """ Args: @@ -73,6 +75,10 @@ def __init__( if set to `None`, use the `num_workers` of the transform data loader. detach: whether to detach the tensors. Scalars tensors will be detached into number types instead of torch tensors. + pad_batch: when the items in a batch indicate different batch size, + whether to pad all the sequences to the longest. + If False, the batch size will be the length of the shortest sequence. + fill_value: the value to fill the padded sequences when `pad_batch=True`. """ self.transform = transform @@ -80,10 +86,12 @@ def __init__( self.num_workers = loader.num_workers if num_workers is None else num_workers self.collate_fn = collate_fn self.detach = detach + self.pad_batch = pad_batch + self.fill_value = fill_value self.pad_collation_used = loader.collate_fn.__doc__ == pad_list_data_collate.__doc__ def __call__(self, data: Dict[str, Any]) -> Any: - decollated_data = decollate_batch(data, detach=self.detach) + decollated_data = decollate_batch(data, detach=self.detach, pad=self.pad_batch, fill_value=self.fill_value) inv_ds = _BatchInverseDataset(decollated_data, self.transform, self.pad_collation_used) inv_loader = DataLoader( inv_ds, batch_size=self.batch_size, num_workers=self.num_workers, collate_fn=self.collate_fn @@ -99,25 +107,36 @@ def __call__(self, data: Dict[str, Any]) -> Any: class Decollated(MapTransform): """ - Decollate a batch of data, if input a dictionary, it can also support to only decollate specified keys. - Note that unlike most MapTransforms, it will delete other keys not specified and if keys=None, will decollate - all the data in the input. - And it replicates the scalar values to every item of the decollated list. + Decollate a batch of data. If input is a dictionary, it also supports to only decollate specified keys. + Note that unlike most MapTransforms, it will delete the other keys that are not specified. + if `keys=None`, it will decollate all the data in the input. + It replicates the scalar values to every item of the decollated list. Args: keys: keys of the corresponding items to decollate, note that it will delete other keys not specified. if None, will decollate all the keys. see also: :py:class:`monai.transforms.compose.MapTransform`. detach: whether to detach the tensors. Scalars tensors will be detached into number types instead of torch tensors. + pad_batch: when the items in a batch indicate different batch size, + whether to pad all the sequences to the longest. + If False, the batch size will be the length of the shortest sequence. + fill_value: the value to fill the padded sequences when `pad_batch=True`. allow_missing_keys: don't raise exception if key is missing. """ def __init__( - self, keys: Optional[KeysCollection] = None, detach: bool = True, allow_missing_keys: bool = False + self, + keys: Optional[KeysCollection] = None, + detach: bool = True, + pad_batch: bool = True, + fill_value=None, + allow_missing_keys: bool = False, ) -> None: super().__init__(keys, allow_missing_keys) self.detach = detach + self.pad_batch = pad_batch + self.fill_value = fill_value def __call__(self, data: Union[Dict, List]): d: Union[Dict, List] @@ -131,4 +150,4 @@ def __call__(self, data: Union[Dict, List]): for key in self.key_iterator(data): d[key] = data[key] - return decollate_batch(rep_scalar_to_batch(d), detach=self.detach) + return decollate_batch(d, detach=self.detach, pad=self.pad_batch, fill_value=self.fill_value) diff --git a/tests/test_decollate.py b/tests/test_decollate.py index 763bf808d0..fc0622c8f0 100644 --- a/tests/test_decollate.py +++ b/tests/test_decollate.py @@ -75,6 +75,7 @@ [[None, None], [None, None]], [["test"], ["test"]], [[], []], + [[("ch1", "ch2"), ("ch3",)], [["ch1", "ch3"], ["ch2", None]]], # default pad None ] @@ -207,6 +208,16 @@ def test_dict_examples(self): out = decollate_batch(test_case, detach=False) self.assertEqual(out[0]["out"], "test") + test_case = { + "image": torch.tensor([[[1, 2, 3]], [[3, 4, 5]]]), + "label": torch.tensor([[[5]], [[7]]]), + "out": ["test"], + } + out = decollate_batch(test_case, detach=False, pad=False) + self.assertEqual(len(out), 1) # no padding + out = decollate_batch(test_case, detach=False, pad=True, fill_value=0) + self.assertEqual(out[1]["out"], 0) # verify padding fill_value + def test_decollated(self): test_case = { "image": torch.tensor([[[1, 2]], [[3, 4]]]), @@ -233,12 +244,16 @@ def test_decollated(self): torch.tensor([[[1, 2]], [[3, 4]]]), {"out": ["test", "test"]}, {"scl_slope": torch.Tensor((0.0, 0.0))}, + {"out2": ["test1"]}, 0.85, + [], ] - transform = Decollated(keys=None, detach=False) + transform = Decollated(keys=None, detach=False, fill_value=-1) out = transform(test_case) - # the 4th item in the list is scalar loss value - self.assertEqual(out[1][3], 0.85) + + self.assertEqual(out[0][-2], 0.85) # scalar replicates + self.assertEqual(out[1][-2], 0.85) # scalar replicates + self.assertEqual(out[1][-3], -1) # fill value for the dictionary item self.assertEqual(out[0][1]["out"], "test") self.assertEqual(out[0][2]["scl_slope"], 0.0) self.assertTrue(isinstance(out[0][2]["scl_slope"], torch.Tensor)) From b03ca4564472bc43881c518198ecc88b176b5d29 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Sat, 20 Nov 2021 14:53:28 +0000 Subject: [PATCH 3/4] fixes batch size Signed-off-by: Wenqi Li --- monai/data/utils.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/monai/data/utils.py b/monai/data/utils.py index 51e908c9ab..c4b49c630c 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -298,10 +298,10 @@ def list_data_collate(batch: Sequence): def _non_zipping_check(data): """ - util used by decollate batch, to identify batch size from the collated data. + utility function used by `decollate_batch`, to identify the largest batch size from the collated data. returns batch_size and the list of non-iterable items. """ - b, non_iterable = 1, [] + batch_size, non_iterable = 0, [] for k, v in data.items() if isinstance(data, Mapping) else enumerate(data): if not isinstance(v, Iterable) or isinstance(v, (str, bytes)) or (isinstance(v, torch.Tensor) and v.ndim == 0): # Not running the usual list decollate here: @@ -309,8 +309,8 @@ def _non_zipping_check(data): # torch.tensor(0) is iterable but iter(torch.tensor(0)) raises TypeError: iteration over a 0-d tensor non_iterable.append(k) elif hasattr(v, "__len__"): - b = max(b, len(v)) - return b, non_iterable + batch_size = max(batch_size, len(v)) + return batch_size, non_iterable def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): @@ -384,10 +384,10 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): return list(out_list) if isinstance(batch, Mapping): _dict_list = {key: decollate_batch(batch[key], detach, pad=pad, fill_value=fill_value) for key in batch} + b, non_iterable = _non_zipping_check(_dict_list) + if b == 0: # all non-iterable + return _dict_list if pad: - b, non_iterable = _non_zipping_check(_dict_list) - if len(non_iterable) == len(_dict_list): - return _dict_list for k in non_iterable: _dict_list[k] = [deepcopy(_dict_list[k]) for _ in range(b)] return [dict(zip(_dict_list, item)) for item in zip_longest(*_dict_list.values(), fillvalue=fill_value)] @@ -395,7 +395,7 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): if isinstance(batch, Iterable): _lists = [decollate_batch(b, detach, pad=pad, fill_value=fill_value) for b in batch] b, non_iterable = _non_zipping_check(_lists) # type: ignore - if len(non_iterable) == len(_lists): # all non-iterable + if b == 0: # all non-iterable return _lists if pad: for k in non_iterable: From 25d3c24ce7dae08f7f4dae0946f9b03810d065ca Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Sun, 21 Nov 2021 13:27:19 +0000 Subject: [PATCH 4/4] update based on comments Signed-off-by: Wenqi Li --- monai/data/utils.py | 53 +++++++++++++++++++++++---------------------- 1 file changed, 27 insertions(+), 26 deletions(-) diff --git a/monai/data/utils.py b/monai/data/utils.py index c4b49c630c..e68dd387ff 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -296,13 +296,21 @@ def list_data_collate(batch: Sequence): raise TypeError(re_str) from re -def _non_zipping_check(data): +def _non_zipping_check(batch_data, detach, pad, fill_value): """ - utility function used by `decollate_batch`, to identify the largest batch size from the collated data. - returns batch_size and the list of non-iterable items. + Utility function based on `decollate_batch`, to identify the largest batch size from the collated data. + returns batch_size, the list of non-iterable items, and the dictionary or list with their items decollated. + + See `decollate_batch` for more details. """ + if isinstance(batch_data, Mapping): + _deco = {key: decollate_batch(batch_data[key], detach, pad=pad, fill_value=fill_value) for key in batch_data} + elif isinstance(batch_data, Iterable): + _deco = [decollate_batch(b, detach, pad=pad, fill_value=fill_value) for b in batch_data] + else: + raise NotImplementedError(f"Unable to de-collate: {batch_data}, type: {type(batch_data)}.") batch_size, non_iterable = 0, [] - for k, v in data.items() if isinstance(data, Mapping) else enumerate(data): + for k, v in _deco.items() if isinstance(_deco, Mapping) else enumerate(_deco): if not isinstance(v, Iterable) or isinstance(v, (str, bytes)) or (isinstance(v, torch.Tensor) and v.ndim == 0): # Not running the usual list decollate here: # don't decollate ['test', 'test'] into [['t', 't'], ['e', 'e'], ['s', 's'], ['t', 't']] @@ -310,7 +318,7 @@ def _non_zipping_check(data): non_iterable.append(k) elif hasattr(v, "__len__"): batch_size = max(batch_size, len(v)) - return batch_size, non_iterable + return batch_size, non_iterable, _deco def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): @@ -367,7 +375,7 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): instead of torch tensors. pad: when the items in a batch indicate different batch size, whether to pad all the sequences to the longest. If False, the batch size will be the length of the shortest sequence. - fill_value: when `pad` is True, the `fillvalue` to use when padding. + fill_value: when `pad` is True, the `fillvalue` to use when padding, defaults to `None`. """ if batch is None: return batch @@ -382,26 +390,19 @@ def decollate_batch(batch, detach: bool = True, pad=True, fill_value=None): 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, pad=pad, fill_value=fill_value) for key in batch} - b, non_iterable = _non_zipping_check(_dict_list) - if b == 0: # all non-iterable - return _dict_list - if pad: - for k in non_iterable: - _dict_list[k] = [deepcopy(_dict_list[k]) for _ in range(b)] - return [dict(zip(_dict_list, item)) for item in zip_longest(*_dict_list.values(), fillvalue=fill_value)] - return [dict(zip(_dict_list, item)) for item in zip(*_dict_list.values())] - if isinstance(batch, Iterable): - _lists = [decollate_batch(b, detach, pad=pad, fill_value=fill_value) for b in batch] - b, non_iterable = _non_zipping_check(_lists) # type: ignore - if b == 0: # all non-iterable - return _lists - if pad: - for k in non_iterable: - _lists[k] = [deepcopy(_lists[k]) for _ in range(b)] - return [list(item) for item in zip_longest(*_lists, fillvalue=fill_value)] - return [list(item) for item in zip(*_lists)] + + b, non_iterable, deco = _non_zipping_check(batch, detach, pad, fill_value) + if b <= 0: # all non-iterable, single item "batch"? {"image": 1, "label": 1} + return deco + if pad: # duplicate non-iterable items to the longest batch + for k in non_iterable: + deco[k] = [deepcopy(deco[k]) for _ in range(b)] + if isinstance(deco, Mapping): + _gen = zip_longest(*deco.values(), fillvalue=fill_value) if pad else zip(*deco.values()) + return [dict(zip(deco, item)) for item in _gen] + if isinstance(deco, Iterable): + _gen = zip_longest(*deco, fillvalue=fill_value) if pad else zip(*deco) + return [list(item) for item in _gen] raise NotImplementedError(f"Unable to de-collate: {batch}, type: {type(batch)}.")