support registering a dataset object; support customizing collate_fn

Reviewed By: zhanghang1989

Differential Revision: D31201334

fbshipit-source-id: 38928aa0eec0749af1edd10f40383257b29bec3c

Author: ppwwyyxx

detectron2/data/build.py

@@ -242,6 +242,9 @@ def get_detection_dataset_dicts(names, filter_empty=True, min_keypoints=0, propo
             for dataset_i_dicts, proposal_file in zip(dataset_dicts, proposal_files)
         ]
 
+    if isinstance(dataset_dicts[0], torchdata.Dataset):
+        return torchdata.ConcatDataset(dataset_dicts)
+
     dataset_dicts = list(itertools.chain.from_iterable(dataset_dicts))
 
     has_instances = "annotations" in dataset_dicts[0]
@@ -263,7 +266,13 @@ def get_detection_dataset_dicts(names, filter_empty=True, min_keypoints=0, propo
 
 
 def build_batch_data_loader(
-    dataset, sampler, total_batch_size, *, aspect_ratio_grouping=False, num_workers=0
+    dataset,
+    sampler,
+    total_batch_size,
+    *,
+    aspect_ratio_grouping=False,
+    num_workers=0,
+    collate_fn=None,
 ):
     """
     Build a batched dataloader. The main differences from `torch.utils.data.DataLoader` are:
@@ -274,7 +283,7 @@ def build_batch_data_loader(
         dataset (torch.utils.data.Dataset): a pytorch map-style or iterable dataset.
         sampler (torch.utils.data.sampler.Sampler or None): a sampler that produces indices.
             Must be provided iff. ``dataset`` is a map-style dataset.
-        total_batch_size, aspect_ratio_grouping, num_workers): see
+        total_batch_size, aspect_ratio_grouping, num_workers, collate_fn: see
             :func:`build_detection_train_loader`.
 
     Returns:
@@ -301,14 +310,17 @@ def build_batch_data_loader(
             collate_fn=operator.itemgetter(0),  # don't batch, but yield individual elements
             worker_init_fn=worker_init_reset_seed,
         )  # yield individual mapped dict
-        return AspectRatioGroupedDataset(data_loader, batch_size)
+        data_loader = AspectRatioGroupedDataset(data_loader, batch_size)
+        if collate_fn is None:
+            return data_loader
+        return MapDataset(data_loader, collate_fn)
     else:
         return torchdata.DataLoader(
             dataset,
             batch_size=batch_size,
             drop_last=True,
             num_workers=num_workers,
-            collate_fn=trivial_batch_collator,
+            collate_fn=trivial_batch_collator if collate_fn is None else collate_fn,
             worker_init_fn=worker_init_reset_seed,
         )
 
@@ -356,7 +368,14 @@ def _train_loader_from_config(cfg, mapper=None, *, dataset=None, sampler=None):
 
 @configurable(from_config=_train_loader_from_config)
 def build_detection_train_loader(
-    dataset, *, mapper, sampler=None, total_batch_size, aspect_ratio_grouping=True, num_workers=0
+    dataset,
+    *,
+    mapper,
+    sampler=None,
+    total_batch_size,
+    aspect_ratio_grouping=True,
+    num_workers=0,
+    collate_fn=None,
 ):
     """
     Build a dataloader for object detection with some default features.
@@ -380,6 +399,11 @@ def build_detection_train_loader(
             aspect ratio for efficiency. When enabled, it requires each
             element in dataset be a dict with keys "width" and "height".
         num_workers (int): number of parallel data loading workers
+        collate_fn: same as the argument of `torch.utils.data.DataLoader`.
+            Defaults to do no collation and return a list of data.
+            No collation is OK for small batch size and simple data structures.
+            If your batch size is large and each sample contains too many small tensors,
+            it's more efficient to collate them in data loader.
 
     Returns:
         torch.utils.data.DataLoader:
@@ -404,6 +428,7 @@ def build_detection_train_loader(
         total_batch_size,
         aspect_ratio_grouping=aspect_ratio_grouping,
         num_workers=num_workers,
+        collate_fn=collate_fn,
     )
 
 
@@ -430,7 +455,7 @@ def _test_loader_from_config(cfg, dataset_name, mapper=None):
 
 
 @configurable(from_config=_test_loader_from_config)
-def build_detection_test_loader(dataset, *, mapper, sampler=None, num_workers=0):
+def build_detection_test_loader(dataset, *, mapper, sampler=None, num_workers=0, collate_fn=None):
     """
     Similar to `build_detection_train_loader`, but uses a batch size of 1,
     and :class:`InferenceSampler`. This sampler coordinates all workers to
@@ -449,6 +474,8 @@ def build_detection_test_loader(dataset, *, mapper, sampler=None, num_workers=0)
             which splits the dataset across all workers. Sampler must be None
             if `dataset` is iterable.
         num_workers (int): number of parallel data loading workers
+        collate_fn: same as the argument of `torch.utils.data.DataLoader`.
+            Defaults to do no collation and return a list of data.
 
     Returns:
         DataLoader: a torch DataLoader, that loads the given detection
@@ -479,7 +506,7 @@ def build_detection_test_loader(dataset, *, mapper, sampler=None, num_workers=0)
         batch_size=1,
         sampler=sampler,
         num_workers=num_workers,
-        collate_fn=trivial_batch_collator,
+        collate_fn=trivial_batch_collator if collate_fn is None else collate_fn,
     )