Refactor multi branch

mmdet/datasets/transforms/wrappers.py

@@ -14,12 +14,62 @@ class MultiBranch(BaseTransform):
     r"""Multiple branch pipeline wrapper.
 
     Generate multiple data-augmented versions of the same image.
+    `MultiBranch` needs to specify the branch names of all
+    pipelines of the dataset, perform corresponding data augmentation
+    for the current branch, and return None for other branches,
+    which ensures the uniformity of return value.
 
     Args:
         branch_field (list): List of branch names.
         branch_pipelines (dict): Dict of different pipeline configs
             to be composed.
 
+    Examples:
+        >>> branch_field = ['sup', 'unsup_teacher', 'unsup_student']
+        >>> sup_pipeline = [
+        >>>     dict(type='LoadImageFromFile',
+        >>>         file_client_args=file_client_args),
+        >>>     dict(type='LoadAnnotations', with_bbox=True),
+        >>>     dict(type='RandomResize', scale=scale, keep_ratio=True),
+        >>>     dict(type='RandomFlip', prob=0.5),
+        >>>     dict(
+        >>>         type='MultiBranch',
+        >>>         branch_field=branch_field,
+        >>>         sup=dict(type='PackDetInputs'))
+        >>>     ]
+        >>> weak_pipeline = [
+        >>>     dict(type='LoadImageFromFile',
+        >>>         file_client_args=file_client_args),
+        >>>     dict(type='LoadAnnotations', with_bbox=True),
+        >>>     dict(type='RandomResize', scale=scale, keep_ratio=True),
+        >>>     dict(type='RandomFlip', prob=0.0),
+        >>>     dict(
+        >>>         type='MultiBranch',
+        >>>         branch_field=branch_field,
+        >>>         sup=dict(type='PackDetInputs'))
+        >>>     ]
+        >>> strong_pipeline = [
+        >>>     dict(type='LoadImageFromFile',
+        >>>         file_client_args=file_client_args),
+        >>>     dict(type='LoadAnnotations', with_bbox=True),
+        >>>     dict(type='RandomResize', scale=scale, keep_ratio=True),
+        >>>     dict(type='RandomFlip', prob=1.0),
+        >>>     dict(
+        >>>         type='MultiBranch',
+        >>>         branch_field=branch_field,
+        >>>         sup=dict(type='PackDetInputs'))
+        >>>     ]
+        >>> unsup_pipeline = [
+        >>>     dict(type='LoadImageFromFile',
+        >>>         file_client_args=file_client_args),
+        >>>     dict(type='LoadEmptyAnnotations'),
+        >>>     dict(
+        >>>         type='MultiBranch',
+        >>>         branch_field=branch_field,
+        >>>         unsup_teacher=weak_pipeline,
+        >>>         unsup_student=strong_pipeline)
+        >>>     ]
+
     """
 
     def __init__(self, branch_field: List[str],

mmdet/models/data_preprocessors/data_preprocessor.py

@@ -368,6 +368,87 @@ def forward(
 class MultiBranchDataPreprocessor(BaseDataPreprocessor):
     """DataPreprocessor wrapper for multi-branch data.
 
+    Take semi-supervised object detection as an example, assume that
+    the ratio of labeled data and unlabeled data in a batch is 1:2,
+    `sup` indicates the branch where the labeled data is augmented,
+    `unsup_teacher` and `unsup_student` indicate the branches where
+    the unlabeled data is augmented by different pipeline. Therefore,
+    the input format of multi-branch data is shown as follows:
+
+    .. code-block:: none
+        {
+            'inputs':
+                {
+                    'sup': [Tensor, None, None],
+                    'unsup_teacher': [None, Tensor, Tensor],
+                    'unsup_student': [None, Tensor, Tensor],
+                },
+            'data_sample':
+                {
+                    'sup': [DetDataSample, None, None],
+                    'unsup_teacher': [None, DetDataSample, DetDataSample],
+                    'unsup_student': [NOne, DetDataSample, DetDataSample],
+                }
+        }
+
+    Filter out branches with a value of None:
+
+    .. code-block:: none
+        {
+            'inputs':
+                {
+                    'sup': [Tensor],
+                    'unsup_teacher': [Tensor, Tensor],
+                    'unsup_student': [Tensor, Tensor],
+                },
+            'data_sample':
+                {
+                    'sup': [DetDataSample],
+                    'unsup_teacher': [DetDataSample, DetDataSample],
+                    'unsup_student': [DetDataSample, DetDataSample],
+                }
+        }
+
+    Group data by branch:
+
+    .. code-block:: none
+        {
+            'sup':
+                {
+                    'inputs': [Tensor]
+                    'data_sample': [DetDataSample, DetDataSample]
+                },
+            'unsup_teacher':
+                {
+                    'inputs': [Tensor, Tensor]
+                    'data_sample': [DetDataSample, DetDataSample]
+                },
+            'unsup_student':
+                {
+                    'inputs': [Tensor, Tensor]
+                    'data_sample': [DetDataSample, DetDataSample]
+                },
+        }
+
+    After preprocessing data from different branches,
+    the multi-branch data needs to be reformatted as:
+
+    .. code-block:: none
+        {
+            'inputs':
+                {
+                    'sup': [Tensor],
+                    'unsup_teacher': [Tensor, Tensor],
+                    'unsup_student': [Tensor, Tensor],
+                },
+            'data_sample':
+                {
+                    'sup': [DetDataSample],
+                    'unsup_teacher': [DetDataSample, DetDataSample],
+                    'unsup_student': [DetDataSample, DetDataSample],
+                }
+        }
+
     Args:
         data_preprocessor (:obj:`ConfigDict` or dict): Config of
             :class:`DetDataPreprocessor` to process the input data.
@@ -390,7 +471,7 @@ def forward(self, data: dict, training: bool = False) -> dict:
 
             - 'inputs' (Dict[str, obj:`torch.Tensor`]): The forward data of
                 models from different branches.
-            - 'data_sample' (Dict[str,obj:`DetDataSample`]): The annotation
+            - 'data_sample' (Dict[str, obj:`DetDataSample`]): The annotation
                 info of the sample from different branches.
         """