Expose data loader parameters in build_batch_data_loader and build_de…

…tection_train_loader.

Summary:
X-link: fairinternal/detectron2#598

Pull Request resolved: #5082

Expose the parameter prefetch_factor, persistent_workers, pin_memory in function build_batch_data_loader and build_detection_train_loader.

Reviewed By: fanzhangmeta

Differential Revision: D46069102

fbshipit-source-id: de47eedea9a61849e60db0181d68a17d1aa1437b

detectron2/data/build.py

@@ -289,6 +289,9 @@ def build_batch_data_loader(
     num_workers=0,
     collate_fn=None,
     drop_last: bool = True,
+    prefetch_factor=None,
+    persistent_workers=False,
+    pin_memory=False,
 ):
     """
     Build a batched dataloader. The main differences from `torch.utils.data.DataLoader` are:
@@ -327,6 +330,9 @@ def build_batch_data_loader(
             num_workers=num_workers,
             collate_fn=operator.itemgetter(0),  # don't batch, but yield individual elements
             worker_init_fn=worker_init_reset_seed,
+            prefetch_factor=prefetch_factor,
+            persistent_workers=persistent_workers,
+            pin_memory=pin_memory,
         )  # yield individual mapped dict
         data_loader = AspectRatioGroupedDataset(data_loader, batch_size)
         if collate_fn is None:
@@ -340,6 +346,9 @@ def build_batch_data_loader(
             num_workers=num_workers,
             collate_fn=trivial_batch_collator if collate_fn is None else collate_fn,
             worker_init_fn=worker_init_reset_seed,
+            prefetch_factor=prefetch_factor,
+            persistent_workers=persistent_workers,
+            pin_memory=pin_memory,
         )
 
 
@@ -479,6 +488,9 @@ def build_detection_train_loader(
     aspect_ratio_grouping=True,
     num_workers=0,
     collate_fn=None,
+    prefetch_factor=None,
+    persistent_workers=False,
+    pin_memory=False,
 ):
     """
     Build a dataloader for object detection with some default features.
@@ -530,6 +542,9 @@ def build_detection_train_loader(
         aspect_ratio_grouping=aspect_ratio_grouping,
         num_workers=num_workers,
         collate_fn=collate_fn,
+        prefetch_factor=prefetch_factor,
+        persistent_workers=persistent_workers,
+        pin_memory=pin_memory,
     )