化繁为简，优化之前基于default_collate的long_text_data_collate为更明确易理解的ser_collate

Author: KevinNuNu

projects/LayoutLMv3/datasets/__init__.py

@@ -1,5 +1,5 @@
 from .transforms import *  # NOQA
-from .utils import long_text_data_collate
+from .utils import ser_collate
 from .xfund_dataset import XFUNDDataset
 
-__all__ = ['XFUNDDataset', 'long_text_data_collate']
+__all__ = ['XFUNDDataset', 'ser_collate']

projects/LayoutLMv3/datasets/utils.py

@@ -1,67 +1,64 @@
-from typing import Any, Mapping, Sequence
+from typing import Dict, Sequence
 
 import torch
 from mmengine.dataset.utils import COLLATE_FUNCTIONS
-from mmengine.structures import BaseDataElement
 
 
 @COLLATE_FUNCTIONS.register_module()
-def long_text_data_collate(data_batch: Sequence, training: bool = True) -> Any:
-    """This code is referenced from
-    ``mmengine.dataset.utils.default_collate``"""
-    data_item = data_batch[0]
-    data_item_type = type(data_item)
+def ser_collate(data_batch: Sequence, training: bool = True) -> Dict:
+    """A collate function designed for SER.
 
-    if isinstance(data_item, (BaseDataElement, str, bytes)):
-        return data_batch
-    elif isinstance(data_item, tuple) and hasattr(data_item, '_fields'):
-        # named_tuple
-        return data_item_type(*(long_text_data_collate(samples, training)
-                                for samples in zip(*data_batch)))
-    elif isinstance(data_item, list):
-        flattened_data_batch = [
-            sub_item for item in data_batch for sub_item in item
-        ]
-        if training:
-            return flattened_data_batch[:len(data_batch)]
-        else:
-            return flattened_data_batch
-    elif isinstance(data_item, Sequence):
-        # check to make sure that the data_itements in batch have
-        # consistent size
-        it = iter(data_batch)
-        data_item_size = len(next(it))
-        if not all(len(data_item) == data_item_size for data_item in it):
-            raise RuntimeError(
-                'each data_itement in list of batch should be of equal size')
-        transposed = list(zip(*data_batch))
+    Args:
+        data_batch (Sequence): Data sampled from dataset.
+        Like:
+            [
+                {
+                    'inputs': {'input_ids': ..., 'bbox': ..., ...},
+                    'data_samples': ['SERDataSample_1']
+                },
+                {
+                    'inputs': {'input_ids': ..., 'bbox': ..., ...},
+                    'data_samples': ['SERDataSample_1', 'SERDataSample_2', ...]
+                },
+                ...
+            ]
+        training (bool): whether training process or not.
 
-        if isinstance(data_item, tuple):
-            return [
-                long_text_data_collate(samples, training)
-                for samples in transposed
-            ]  # Compat with Pytorch.
-        else:
-            try:
-                return data_item_type([
-                    long_text_data_collate(samples, training)
-                    for samples in transposed
-                ])
-            except TypeError:
-                # The sequence type may not support `__init__(iterable)`
-                # (e.g., `range`).
-                return [
-                    long_text_data_collate(samples, training)
-                    for samples in transposed
-                ]
-    elif isinstance(data_item, Mapping):
-        return data_item_type({
-            key: long_text_data_collate([d[key] for d in data_batch], training)
-            for key in data_item
-        })
-    else:
-        concat_data_batch = torch.concat(data_batch, dim=0)
-        if training:
-            return concat_data_batch[:len(data_batch)]
-        else:
-            return concat_data_batch
+    Note:
+        Different from ``default_collate`` in pytorch or in mmengine,
+        ``ser_collate`` can accept `inputs` tensor and `data_samples`
+        list with the different shape.
+
+    Returns:
+        transposed (Dict): A dict have two elements,
+            the first element `inputs` is a dict
+            the second element `data_samples` is a list
+    """
+    batch_size = len(data_batch)
+    # transpose `inputs`, which is a dict.
+    batch_inputs = [data_item['inputs'] for data_item in data_batch]
+    batch_inputs_item = batch_inputs[0]
+    transposed_batch_inputs = {}
+    for key in batch_inputs_item:
+        concat_value = torch.concat([d[key] for d in batch_inputs], dim=0)
+        # TODO: because long text will be truncated, the concat_value
+        # cannot be sliced directly when training=False.
+        # How to support batch inference?
+        transposed_batch_inputs[key] = concat_value[:batch_size] \
+            if training else concat_value
+    # transpose `data_samples`, which is a list.
+    batch_data_samples = [
+        data_item['data_samples'] for data_item in data_batch
+    ]
+    flattened = [sub_item for item in batch_data_samples for sub_item in item]
+    # TODO: because long text will be truncated, the concat_value
+    # cannot be sliced directly when training=False.
+    # How to support batch inference?
+    transposed_batch_data_samples = flattened[:batch_size] \
+        if training else flattened
+
+    transposed = {
+        'inputs': transposed_batch_inputs,
+        'data_samples': transposed_batch_data_samples
+    }
+    return transposed