add __len__ for ToIterableDataset

Reviewed By: zhanghang1989

Differential Revision: D30635073

fbshipit-source-id: 664ba17b768a69fed97ca94b944e43d186336966

Author: ppwwyyxx

detectron2/data/common.py

@@ -200,6 +200,9 @@ def __iter__(self):
         for idx in sampler:
             yield self.dataset[idx]
 
+    def __len__(self):
+        return len(self.sampler)
+
 
 class AspectRatioGroupedDataset(data.IterableDataset):
     """

tests/data/test_dataset.py

@@ -14,9 +14,10 @@
     DatasetFromList,
     MapDataset,
     ToIterableDataset,
+    build_batch_data_loader,
     build_detection_train_loader,
 )
-from detectron2.data.samplers import TrainingSampler
+from detectron2.data.samplers import InferenceSampler, TrainingSampler
 
 
 def _a_slow_func(x):
@@ -92,3 +93,14 @@ def test_build_iterable_dataloader(self):
         ds = ToIterableDataset(ds, TrainingSampler(len(ds)))
         dl = build_detection_train_loader(dataset=ds, **kwargs)
         next(iter(dl))
+
+    def test_build_dataloader_inference(self):
+        N = 96
+        ds = DatasetFromList(list(range(N)))
+        sampler = InferenceSampler(len(ds))
+        dl = build_batch_data_loader(ds, sampler, 8, num_workers=3)
+
+        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)))