[Feature] Support custom persistent_workers (#6435)

* Fix aug test error when the number of prediction bboxes is 0 (#6398)

* Fix aug test error when the number of prediction bboxes is 0

* test

* test

* fix lint

* Support custom pin_memory and persistent_workers

* fix comment

* fix docstr

* remove pin_memory

mmdet/apis/train.py

@@ -105,7 +105,9 @@ def train_detector(model,
             num_gpus=len(cfg.gpu_ids),
             dist=distributed,
             seed=cfg.seed,
-            runner_type=runner_type) for ds in dataset
+            runner_type=runner_type,
+            persistent_workers=cfg.data.get('persistent_workers', False))
+        for ds in dataset
     ]
 
     # put model on gpus

mmdet/datasets/builder.py

@@ -2,12 +2,13 @@
 import copy
 import platform
 import random
+import warnings
 from functools import partial
 
 import numpy as np
 from mmcv.parallel import collate
 from mmcv.runner import get_dist_info
-from mmcv.utils import Registry, build_from_cfg
+from mmcv.utils import TORCH_VERSION, Registry, build_from_cfg, digit_version
 from torch.utils.data import DataLoader
 
 from .samplers import (DistributedGroupSampler, DistributedSampler,
@@ -90,6 +91,7 @@ def build_dataloader(dataset,
                      shuffle=True,
                      seed=None,
                      runner_type='EpochBasedRunner',
+                     persistent_workers=False,
                      **kwargs):
     """Build PyTorch DataLoader.
 
@@ -106,7 +108,12 @@ def build_dataloader(dataset,
         dist (bool): Distributed training/test or not. Default: True.
         shuffle (bool): Whether to shuffle the data at every epoch.
             Default: True.
+        seed (int, Optional): Seed to be used. Default: None.
         runner_type (str): Type of runner. Default: `EpochBasedRunner`
+        persistent_workers (bool): If True, the data loader will not shutdown
+            the worker processes after a dataset has been consumed once.
+            This allows to maintain the workers `Dataset` instances alive.
+            This argument is only valid when PyTorch>=1.7.0. Default: False.
         kwargs: any keyword argument to be used to initialize DataLoader
 
     Returns:
@@ -163,6 +170,13 @@ def build_dataloader(dataset,
         worker_init_fn, num_workers=num_workers, rank=rank,
         seed=seed) if seed is not None else None
 
+    if (TORCH_VERSION != 'parrots'
+            and digit_version(TORCH_VERSION) >= digit_version('1.7.0')):
+        kwargs['persistent_workers'] = persistent_workers
+    elif persistent_workers is True:
+        warnings.warn('persistent_workers is invalid because your pytorch '
+                      'version is lower than 1.7.0')
+
     data_loader = DataLoader(
         dataset,
         batch_size=batch_size,