Repository: Project-MONAI/MONAI

Issue:
simplify `list_data_collate` and `collate_meta_tensor ` using `collate`
the generic `collate` API is [implemented in pytorch ](https://github.com/pytorch/pytorch/blob/d7a3f2128fb4457dd60fd5d23e77d2c66a8b0f02/torch/utils/data/_utils/collate.py#L88-L92)

https://github.com/pytorch/pytorch/pull/85748

```
def collate(batch, *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None):
    r"""
        General collate function that handles collection type of element within each batch
        and opens function registry to deal with specific element types. `default_collate_fn_map`
        provides default collate functions for tensors, numpy arrays, numbers and strings.

```

probably the monai helper functions could be simplified:
https://github.com/Project-MONAI/MONAI/blob/50c3f328cb98e4f0c4e6a38a04a98badf0961c18/monai/data/utils.py#L448-L451

https://github.com/Project-MONAI/MONAI/blob/50c3f328cb98e4f0c4e6a38a04a98badf0961c18/monai/data/utils.py#L427-L430
simplify `list_data_collate` and `collate_meta_tensor ` using `collate`
the generic `collate` API is [implemented in pytorch ](https://github.com/pytorch/pytorch/blob/d7a3f2128fb4457dd60fd5d23e77d2c66a8b0f02/torch/utils/data/_utils/collate.py#L88-L92)

https://github.com/pytorch/pytorch/pull/85748

```
def collate(batch, *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None):
    r"""
        General collate function that handles collection type of element within each batch
        and opens function registry to deal with specific element types. `default_collate_fn_map`
        provides default collate functions for tensors, numpy arrays, numbers and strings.

```

probably the monai helper functions could be simplified:
https://github.com/Project-MONAI/MONAI/blob/50c3f328cb98e4f0c4e6a38a04a98badf0961c18/monai/data/utils.py#L448-L451

https://github.com/Project-MONAI/MONAI/blob/50c3f328cb98e4f0c4e6a38a04a98badf0961c18/monai/data/utils.py#L427-L430


Discussion:
It may also affect this feature request https://github.com/Project-MONAI/MONAI/discussions/5909?
CC @yiheng-wang-nv 
It would be cool if in this contribution, the `collate_meta_tensor` could be more permissive, if keys in the metadata are missing. I work with data from multiple datasets and frequently get errors, as keys are not the same across the different images. I usually just pop these keys in the metadata, but it would be nice, if this would not be an error at all. 

```
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/worker.py", line 302, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 61, in fetch
    return self.collate_fn(data)
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/monai/data/utils.py", line 470, in list_data_collate
    ret = collate_meta_tensor(data)
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/monai/data/utils.py", line 439, in collate_meta_tensor
    return {k: collate_meta_tensor([d[k] for d in batch]) for k in elem_0}
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/monai/data/utils.py", line 439, in <dictcomp>
    return {k: collate_meta_tensor([d[k] for d in batch]) for k in elem_0}
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/monai/data/utils.py", line 434, in collate_meta_tensor
    collated.meta = default_collate([i.meta or TraceKeys.NONE for i in batch])
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/collate.py", line 265, in default_collate
    return collate(batch, collate_fn_map=default_collate_fn_map)
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/collate.py", line 128, in collate
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/collate.py", line 128, in <dictcomp>
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/collate.py", line 128, in <listcomp>
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
KeyError: 'ITK_non_uniform_sampling_deviation'
```
Hi @wyli ,

Could you please help double confirm it?

Thanks in advance.

created an additional feature request to make sure the requirement tracked
It may also affect this feature request https://github.com/Project-MONAI/MONAI/discussions/5909?
CC @yiheng-wang-nv 
It would be cool if in this contribution, the `collate_meta_tensor` could be more permissive, if keys in the metadata are missing. I work with data from multiple datasets and frequently get errors, as keys are not the same across the different images. I usually just pop these keys in the metadata, but it would be nice, if this would not be an error at all. 

```
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/worker.py", line 302, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 61, in fetch
    return self.collate_fn(data)
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/monai/data/utils.py", line 470, in list_data_collate
    ret = collate_meta_tensor(data)
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/monai/data/utils.py", line 439, in collate_meta_tensor
    return {k: collate_meta_tensor([d[k] for d in batch]) for k in elem_0}
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/monai/data/utils.py", line 439, in <dictcomp>
    return {k: collate_meta_tensor([d[k] for d in batch]) for k in elem_0}
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/monai/data/utils.py", line 434, in collate_meta_tensor
    collated.meta = default_collate([i.meta or TraceKeys.NONE for i in batch])
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/collate.py", line 265, in default_collate
    return collate(batch, collate_fn_map=default_collate_fn_map)
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/collate.py", line 128, in collate
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/collate.py", line 128, in <dictcomp>
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
  File "/home/bressekk/miniconda3/envs/monailabel/lib/python3.9/site-packages/torch/utils/data/_utils/collate.py", line 128, in <listcomp>
    return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
KeyError: 'ITK_non_uniform_sampling_deviation'
```
Hi @wyli ,

Could you please help double confirm it?

Thanks in advance.

created an additional feature request to make sure the requirement tracked

Resolving patch:
diff --git a/monai/data/utils.py b/monai/data/utils.py
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -50,8 +50,12 @@
     issequenceiterable,
     look_up_option,
     optional_import,
+    pytorch_after,
 )
 
+if pytorch_after(1, 13):
+    # import private code for reuse purposes, comment in case things break in the future
+    from torch.utils.data._utils.collate import collate_tensor_fn, default_collate_fn_map
 pd, _ = optional_import("pandas")
 DataFrame, _ = optional_import("pandas", name="DataFrame")
 nib, _ = optional_import("nibabel")
@@ -444,6 +448,23 @@ def pickle_operations(data, key=PICKLE_KEY_SUFFIX, is_encode: bool = True):
     return data
 
 
+def collate_meta_tensor_fn(batch, *, collate_fn_map=None):
+    """
+    Collate a sequence of meta tensor into a single batched metatensor. This is called by `collage_meta_tensor`
+    and so should not be used as a collate function directly in dataloaders.
+    """
+    collate_fn = collate_tensor_fn if pytorch_after(1, 13) else default_collate
+    collated = collate_fn(batch)  # type: ignore
+    meta_dicts = [i.meta or TraceKeys.NONE for i in batch]
+    common_ = set.intersection(*[set(d.keys()) for d in meta_dicts if isinstance(d, dict)])
+    if common_:
+        meta_dicts = [{k: d[k] for k in common_} if isinstance(d, dict) else TraceKeys.NONE for d in meta_dicts]
+    collated.meta = default_collate(meta_dicts)
+    collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]
+    collated.is_batch = True
+    return collated
+
+
 def collate_meta_tensor(batch):
     """collate a sequence of meta tensor sequences/dictionaries into
     a single batched metatensor or a dictionary of batched metatensor"""
@@ -451,15 +472,7 @@ def collate_meta_tensor(batch):
         raise NotImplementedError()
     elem_0 = first(batch)
     if isinstance(elem_0, MetaObj):
-        collated = default_collate(batch)
-        meta_dicts = [i.meta or TraceKeys.NONE for i in batch]
-        common_ = set.intersection(*[set(d.keys()) for d in meta_dicts if isinstance(d, dict)])
-        if common_:
-            meta_dicts = [{k: d[k] for k in common_} if isinstance(d, dict) else TraceKeys.NONE for d in meta_dicts]
-        collated.meta = default_collate(meta_dicts)
-        collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]
-        collated.is_batch = True
-        return collated
+        return collate_meta_tensor_fn(batch)
     if isinstance(elem_0, Mapping):
         return {k: collate_meta_tensor([d[k] for d in batch]) for k in elem_0}
     if isinstance(elem_0, (tuple, list)):
@@ -479,9 +492,16 @@ def list_data_collate(batch: Sequence):
         Need to use this collate if apply some transforms that can generate batch data.
 
     """
+
+    if pytorch_after(1, 13):
+        # needs to go here to avoid circular import
+        from monai.data.meta_tensor import MetaTensor
+
+        default_collate_fn_map.update({MetaTensor: collate_meta_tensor_fn})
     elem = batch[0]
     data = [i for k in batch for i in k] if isinstance(elem, list) else batch
     key = None
+    collate_fn = default_collate if pytorch_after(1, 13) else collate_meta_tensor
     try:
         if config.USE_META_DICT:
             data = pickle_operations(data)  # bc 0.9.0
@@ -490,9 +510,9 @@ def list_data_collate(batch: Sequence):
             for k in elem:
                 key = k
                 data_for_batch = [d[key] for d in data]
-                ret[key] = collate_meta_tensor(data_for_batch)
+                ret[key] = collate_fn(data_for_batch)
         else:
-            ret = collate_meta_tensor(data)
+            ret = collate_fn(data)
         return ret
     except RuntimeError as re:
         re_str = str(re)


Test patch: