Problem with collation when conditional transform emits different structure

[kretes]

I stumbled upon a case when using a custom data transform pipeline.
Background:
I am using a custom strategy of sampling data, where for some data points I use RandCropByPosNegLabeld while for some - RandSpatialCropd
the problem is that RandCropByPosNegLabeld adds some new fields to metadata ('patch_index' and 'crop_center'), that I don't get from RandSpatialCropd.

And so I get errors from collate mechanism, that seems to be hierarchical, and strict in terms of the structure.

The traceback

Original Traceback (most recent call last):
  File "***/lib/python3.11/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
           ^^^^^^^^^^^^^^^^^^^^
  File "***/lib/python3.11/site-packages/torch/utils/data/_utils/fetch.py", line 54, in fetch
    return self.collate_fn(data)
           ^^^^^^^^^^^^^^^^^^^^^
  File "***/lib/python3.11/site-packages/monai/data/utils.py", line 488, in list_data_collate
    ret[key] = collate_meta_tensor(data_for_batch)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "***/lib/python3.11/site-packages/monai/data/utils.py", line 454, in collate_meta_tensor
    collated.meta = default_collate([i.meta or TraceKeys.NONE for i in batch])
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "***/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py", line 265, in default_collate
    return collate(batch, collate_fn_map=default_collate_fn_map)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "***/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py", line 127, in collate
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "***/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py", line 127, in <dictcomp>
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
                                   ^^^^^^^^^^^^^^^^^^^^^^^
  File "***/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py", line 127, in <listcomp>
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
                                    ~^^^^^
KeyError: 'patch_index'

It wasn't very obvious for me at first what's going on and I spent some time to debug this issue.
The problem gets more complicated, as collate mechanism takes the first element from the batch to infer and expect the structure and so in some batches it worked ok, while in other it didn't.

I am not sure about the solution for this, but I think monai should offer something more.
First of all I am thinking about some decent structure checking, e.g. to show the hierarchical aspect of this situation.

It would be nice to get an exception like 'Different structures in batch elements. batch element 0 has image.meta.patch_index, while batch element 2 is missing'

Additional thing is that RandCropByPosNegLabeld doesn't say explicitly in the doc that such keys will be added, and so the user needs to dig through source code / check empirically.

I have the workaround, which is to add at least the same structure in the other branch of code:

new_img_dict["image"].meta["patch_index"] = 0
new_img_dict["image"].meta["crop_center"] = (0, 0, 0)

but this seems hairy and will not generalize into similar cases

[KumoLiu]

Hi @kretes, could you please share a small script I can help debug? I can't reproduce using the code below.

from monai.transforms import RandSpatialCropd, RandCropByPosNegLabeld, LoadImaged, Compose
from monai.data.utils import collate_meta_tensor

data = [
    {
    "image": image_path1,
    "label": label_path1,
    },
    {
    "image": image_path2,
    "label": label_path2,
    },
]

trans1 = Compose([
    LoadImaged(keys=["image", "label"], image_only=True, ensure_channel_first=True),
    RandSpatialCropd(keys=["image", "label"], roi_size=[96,96,3]),
])
trans2 = Compose([
    LoadImaged(keys=["image", "label"], image_only=True, ensure_channel_first=True),
    RandCropByPosNegLabeld(keys=["image", "label"], label_key="label", spatial_size=[96,96,3]),
])

a = trans1(data[0])
b = trans2(data[1])

c = collate_meta_tensor([a["image"], b[0]["image"]])

Thanks!

[kretes]

Yes, your script works (although I needed to add , random_size=False in RandSpatialCropd.

However, if you switch the order of tensors that are passed to collate_meta_tensor -> it fails:

d = collate_meta_tensor([b[0]["image"], a["image"]])

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[17], line 30
     27 b = trans2(data[1])
     29 c = collate_meta_tensor([a["image"], b[0]["image"]])
---> 30 d = collate_meta_tensor([b[0]["image"], a["image"]])

File ***/lib/python3.11/site-packages/monai/data/utils.py:454, in collate_meta_tensor(batch)
    452 if isinstance(elem_0, MetaObj):
    453     collated = default_collate(batch)
--> 454     collated.meta = default_collate([i.meta or TraceKeys.NONE for i in batch])
    455     collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]
    456     collated.is_batch = True

File ***/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py:265, in default_collate(batch)
    204 def default_collate(batch):
    205     r"""
    206         Function that takes in a batch of data and puts the elements within the batch
    207         into a tensor with an additional outer dimension - batch size. The exact output type can be
   (...)
    263             >>> default_collate(batch)  # Handle `CustomType` automatically
    264     """
--> 265     return collate(batch, collate_fn_map=default_collate_fn_map)

File ***/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py:127, in collate(batch, collate_fn_map)
    125 if isinstance(elem, collections.abc.Mapping):
    126     try:
--> 127         return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
    128     except TypeError:
    129         # The mapping type may not support `__init__(iterable)`.
    130         return {key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem}

File ***/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py:127, in <dictcomp>(.0)
    125 if isinstance(elem, collections.abc.Mapping):
    126     try:
--> 127         return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
    128     except TypeError:
    129         # The mapping type may not support `__init__(iterable)`.
    130         return {key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem}

File ***/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py:127, in <listcomp>(.0)
    125 if isinstance(elem, collections.abc.Mapping):
    126     try:
--> 127         return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
    128     except TypeError:
    129         # The mapping type may not support `__init__(iterable)`.
    130         return {key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem}

KeyError: 'patch_index'

Which is as well an API glitch, since it should be equivariant .
Why this happens is somewhere inside collate - the first element is used as a reference and the keys are taken from that one.

[KumoLiu]

oh, I see. Perhaps you could provide a customized collate_fn in DataLoader such as skipping this key or generating some fake data. Thanks!

[wyli]

also instead of changing collate_fn, do you think enhancing RandSpatialCropSamples can address this issue? the transform seems just miss a meta["crop_center"]

MONAI/monai/transforms/croppad/array.py

Lines 701 to 705 in fabc198

	class RandSpatialCropSamples(Randomizable, TraceableTransform, LazyTransform, MultiSampleTrait):
	"""
	Crop image with random size or specific size ROI to generate a list of N samples.
	It can crop at a random position as center or at the image center. And allows to set
	the minimum size to limit the randomly generated ROI.

[kretes]

It seems to miss both crop_center and patch_index.
Yes, filling those in would be sensible.

My discussion is as well about the entangled way this error is shown to the user. But I don't have any idea about it. collation seems quite complex operation to handle a lot of edge cases anyway

[wyli]

thanks for raising the issue, we added some logic to identify common meta dictionary keys. #7054