support iterable dataset in build_detection_test_loader

Summary: Like D24677397 (05bc843), but for test loader

Differential Revision: D31161853

fbshipit-source-id: 6c101843a2be681fc23b2ff241070876a77be80f

Author: ppwwyyxx

detectron2/data/build.py

@@ -439,14 +439,15 @@ def build_detection_test_loader(dataset, *, mapper, sampler=None, num_workers=0)
 
     Args:
         dataset (list or torch.utils.data.Dataset): a list of dataset dicts,
-            or a map-style pytorch dataset. They can be obtained by using
-            :func:`DatasetCatalog.get` or :func:`get_detection_dataset_dicts`.
+            or a pytorch dataset (either map-style or iterable). They can be obtained
+            by using :func:`DatasetCatalog.get` or :func:`get_detection_dataset_dicts`.
         mapper (callable): a callable which takes a sample (dict) from dataset
            and returns the format to be consumed by the model.
            When using cfg, the default choice is ``DatasetMapper(cfg, is_train=False)``.
         sampler (torch.utils.data.sampler.Sampler or None): a sampler that produces
             indices to be applied on ``dataset``. Default to :class:`InferenceSampler`,
-            which splits the dataset across all workers.
+            which splits the dataset across all workers. Sampler must be None
+            if `dataset` is iterable.
         num_workers (int): number of parallel data loading workers
 
     Returns:
@@ -466,18 +467,20 @@ def build_detection_test_loader(dataset, *, mapper, sampler=None, num_workers=0)
         dataset = DatasetFromList(dataset, copy=False)
     if mapper is not None:
         dataset = MapDataset(dataset, mapper)
-    if sampler is None:
-        sampler = InferenceSampler(len(dataset))
+    if isinstance(dataset, torchdata.IterableDataset):
+        assert sampler is None, "sampler must be None if dataset is IterableDataset"
+    else:
+        if sampler is None:
+            sampler = InferenceSampler(len(dataset))
     # Always use 1 image per worker during inference since this is the
     # standard when reporting inference time in papers.
-    batch_sampler = torchdata.sampler.BatchSampler(sampler, 1, drop_last=False)
-    data_loader = torchdata.DataLoader(
+    return torchdata.DataLoader(
         dataset,
+        batch_size=1,
+        sampler=sampler,
         num_workers=num_workers,
-        batch_sampler=batch_sampler,
         collate_fn=trivial_batch_collator,
     )
-    return data_loader
 
 
 def trivial_batch_collator(batch):

tests/data/test_dataset.py

@@ -15,6 +15,7 @@
     MapDataset,
     ToIterableDataset,
     build_batch_data_loader,
+    build_detection_test_loader,
     build_detection_train_loader,
 )
 from detectron2.data.samplers import InferenceSampler, TrainingSampler
@@ -82,25 +83,46 @@ def _get_kwargs(self):
         kwargs = {k: instantiate(v) for k, v in cfg.items()}
         return kwargs
 
-    def test_build_dataloader(self):
+    def test_build_dataloader_train(self):
         kwargs = self._get_kwargs()
         dl = build_detection_train_loader(**kwargs)
         next(iter(dl))
 
-    def test_build_iterable_dataloader(self):
+    def test_build_iterable_dataloader_train(self):
         kwargs = self._get_kwargs()
         ds = DatasetFromList(kwargs.pop("dataset"))
         ds = ToIterableDataset(ds, TrainingSampler(len(ds)))
         dl = build_detection_train_loader(dataset=ds, **kwargs)
         next(iter(dl))
 
-    def test_build_dataloader_inference(self):
+    def _check_is_range(self, data_loader, N):
+        # check that data_loader produces range(N)
+        data = list(iter(data_loader))
+        data = [x for batch in data for x in batch]  # flatten the batches
+        self.assertEqual(len(data), N)
+        self.assertEqual(set(data), set(range(N)))
+
+    def test_build_batch_dataloader_inference(self):
+        # Test that build_batch_data_loader can be used for inference
         N = 96
         ds = DatasetFromList(list(range(N)))
         sampler = InferenceSampler(len(ds))
         dl = build_batch_data_loader(ds, sampler, 8, num_workers=3)
+        self._check_is_range(dl, N)
 
-        data = list(iter(dl))
-        data = [x for batch in data for x in batch]  # flatten the batches
-        self.assertEqual(len(data), N)
-        self.assertEqual(set(data), set(range(N)))
+    def test_build_dataloader_inference(self):
+        N = 50
+        ds = DatasetFromList(list(range(N)))
+        sampler = InferenceSampler(len(ds))
+        dl = build_detection_test_loader(
+            dataset=ds, sampler=sampler, mapper=lambda x: x, num_workers=3
+        )
+        self._check_is_range(dl, N)
+
+    def test_build_iterable_dataloader_inference(self):
+        # Test that build_detection_test_loader supports iterable dataset
+        N = 50
+        ds = DatasetFromList(list(range(N)))
+        ds = ToIterableDataset(ds, InferenceSampler(len(ds)))
+        dl = build_detection_test_loader(dataset=ds, mapper=lambda x: x, num_workers=3)
+        self._check_is_range(dl, N)