use DataLoader for testing on select configs

Author: pmeier

test/datasets_utils.py

@@ -3,9 +3,9 @@
 import importlib
 import inspect
 import itertools
+import multiprocessing
 import os
 import pathlib
-import pickle
 import random
 import shutil
 import string
@@ -171,6 +171,38 @@ def wrapper(self):
     return wrapper
 
 
+def _no_collate(batch):
+    return batch
+
+
+def check_transforms_v2_wrapper(dataset_test_case, *, config=None, supports_target_keys=False):
+    from torch.utils.data import DataLoader
+    from torchvision import datapoints
+    from torchvision.datasets import wrap_dataset_for_transforms_v2
+
+    target_keyss = [None]
+    if supports_target_keys:
+        target_keyss.append("all")
+
+    for target_keys, multiprocessing_context in itertools.product(
+        target_keyss, multiprocessing.get_all_start_methods()
+    ):
+        with dataset_test_case.create_dataset(config) as (dataset, info):
+            wrapped_dataset = wrap_dataset_for_transforms_v2(dataset, target_keys=target_keys)
+
+            assert isinstance(wrapped_dataset, type(dataset))
+            assert len(wrapped_dataset) == info["num_examples"]
+
+            dataloader = DataLoader(
+                wrapped_dataset, num_workers=2, multiprocessing_context=multiprocessing_context, collate_fn=_no_collate
+            )
+
+            for wrapped_sample in dataloader:
+                assert tree_any(
+                    lambda item: isinstance(item, (datapoints.Image, datapoints.Video, PIL.Image.Image)), wrapped_sample
+                )
+
+
 class DatasetTestCase(unittest.TestCase):
     """Abstract base class for all dataset testcases.
 
@@ -566,49 +598,6 @@ def test_transforms(self, config):
 
                 mock.assert_called()
 
-    @test_all_configs
-    def test_transforms_v2_wrapper(self, config):
-        from torchvision import datapoints
-        from torchvision.datasets import wrap_dataset_for_transforms_v2
-
-        try:
-            with self.create_dataset(config) as (dataset, info):
-                wrap_dataset_for_transforms_v2(dataset)
-        except TypeError as error:
-            msg = f"No wrapper exists for dataset class {type(dataset).__name__}"
-            if str(error).startswith(msg):
-                return
-            raise error
-        except RuntimeError as error:
-            if "currently not supported by this wrapper" in str(error):
-                return
-            raise error
-
-        for target_keys, de_serialize in itertools.product(
-            [None, "all"], [lambda d: d, lambda d: pickle.loads(pickle.dumps(d))]
-        ):
-
-            with self.create_dataset(config) as (dataset, info):
-                if target_keys is not None and self.DATASET_CLASS not in {
-                    torchvision.datasets.CocoDetection,
-                    torchvision.datasets.VOCDetection,
-                    torchvision.datasets.Kitti,
-                    torchvision.datasets.WIDERFace,
-                }:
-                    with self.assertRaisesRegex(ValueError, "`target_keys` is currently only supported for"):
-                        wrap_dataset_for_transforms_v2(dataset, target_keys=target_keys)
-                    continue
-
-                wrapped_dataset = de_serialize(wrap_dataset_for_transforms_v2(dataset, target_keys=target_keys))
-
-                assert isinstance(wrapped_dataset, self.DATASET_CLASS)
-                assert len(wrapped_dataset) == info["num_examples"]
-
-                wrapped_sample = wrapped_dataset[0]
-                assert tree_any(
-                    lambda item: isinstance(item, (datapoints.Image, datapoints.Video, PIL.Image.Image)), wrapped_sample
-                )
-
 
 class ImageDatasetTestCase(DatasetTestCase):
     """Abstract base class for image dataset testcases.
@@ -690,15 +679,6 @@ def wrapper(tmpdir, config):
 
         return wrapper
 
-    @test_all_configs
-    def test_transforms_v2_wrapper(self, config):
-        # `output_format == "THWC"` is not supported by the wrapper. Thus, we skip the `config` if it is set explicitly
-        # or use the supported `"TCHW"`
-        if config.setdefault("output_format", "TCHW") == "THWC":
-            return
-
-        super().test_transforms_v2_wrapper.__wrapped__(self, config)
-
 
 def create_image_or_video_tensor(size: Sequence[int]) -> torch.Tensor:
     r"""Create a random uint8 tensor.

test/test_datasets.py

@@ -183,6 +183,9 @@ def test_combined_targets(self):
                 ), "Type of the combined target does not match the type of the corresponding individual target: "
                 f"{actual} is not {expected}",
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self, config=dict(target_type="category"))
+
 
 class Caltech256TestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.Caltech256
@@ -203,6 +206,9 @@ def inject_fake_data(self, tmpdir, config):
 
         return num_images_per_category * len(categories)
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class WIDERFaceTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.WIDERFace
@@ -258,6 +264,9 @@ def inject_fake_data(self, tmpdir, config):
 
         return split_to_num_examples[config["split"]]
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self, supports_target_keys=True)
+
 
 class CityScapesTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.Cityscapes
@@ -382,6 +391,10 @@ def test_feature_types_target_polygon(self):
             assert isinstance(polygon_img, PIL.Image.Image)
             (polygon_target, info["expected_polygon_target"])
 
+    def test_transforms_v2_wrapper(self):
+        for target_type in ["instance", "semantic", ["instance", "semantic"]]:
+            datasets_utils.check_transforms_v2_wrapper(self, config=dict(target_type=target_type))
+
 
 class ImageNetTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.ImageNet
@@ -413,6 +426,9 @@ def inject_fake_data(self, tmpdir, config):
         torch.save((wnid_to_classes, None), tmpdir / "meta.bin")
         return num_examples
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.CIFAR10
@@ -470,6 +486,9 @@ def test_class_to_idx(self):
             actual = dataset.class_to_idx
             assert actual == expected
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class CIFAR100(CIFAR10TestCase):
     DATASET_CLASS = datasets.CIFAR100
@@ -484,6 +503,9 @@ class CIFAR100(CIFAR10TestCase):
         categories_key="fine_label_names",
     )
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class CelebATestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.CelebA
@@ -607,6 +629,10 @@ def test_images_names_split(self):
 
         assert merged_imgs_names == all_imgs_names
 
+    def test_transforms_v2_wrapper(self):
+        for target_type in ["identity", "bbox", ["identity", "bbox"]]:
+            datasets_utils.check_transforms_v2_wrapper(self, config=dict(target_type=target_type))
+
 
 class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.VOCSegmentation
@@ -694,6 +720,9 @@ def add_bndbox(obj, bndbox=None):
 
         return data
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class VOCDetectionTestCase(VOCSegmentationTestCase):
     DATASET_CLASS = datasets.VOCDetection
@@ -714,6 +743,10 @@ def test_annotations(self):
 
             assert object == info["annotation"]
 
+    def test_transforms_v2_wrapper(self):
+        for target_type in ["identity", "bbox", ["identity", "bbox"]]:
+            datasets_utils.check_transforms_v2_wrapper(self, supports_target_keys=True)
+
 
 class CocoDetectionTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.CocoDetection
@@ -784,6 +817,9 @@ def _create_json(self, root, name, content):
             json.dump(content, fh)
         return file
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self, supports_target_keys=True)
+
 
 class CocoCaptionsTestCase(CocoDetectionTestCase):
     DATASET_CLASS = datasets.CocoCaptions
@@ -800,6 +836,11 @@ def test_captions(self):
             _, captions = dataset[0]
             assert tuple(captions) == tuple(info["captions"])
 
+    def test_transforms_v2_wrapper(self):
+        # We need to define this method, because otherwise the test from the super class will
+        # be run
+        pytest.skip("CocoCaptions is currently not supported by the v2 wrapper.")
+
 
 class UCF101TestCase(datasets_utils.VideoDatasetTestCase):
     DATASET_CLASS = datasets.UCF101
@@ -860,6 +901,9 @@ def _create_annotation_file(self, root, name, video_files):
         with open(pathlib.Path(root) / name, "w") as fh:
             fh.writelines(f"{str(file).replace(os.sep, '/')}\n" for file in sorted(video_files))
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self, config=dict(output_format="TCHW"))
+
 
 class LSUNTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.LSUN
@@ -966,6 +1010,9 @@ def inject_fake_data(self, tmpdir, config):
             )
         return num_videos_per_class * len(classes)
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self, config=dict(output_format="TCHW"))
+
 
 class HMDB51TestCase(datasets_utils.VideoDatasetTestCase):
     DATASET_CLASS = datasets.HMDB51
@@ -1026,6 +1073,9 @@ def _create_split_files(self, root, video_files, fold, train):
 
         return num_train_videos if train else (num_videos - num_train_videos)
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self, config=dict(output_format="TCHW"))
+
 
 class OmniglotTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.Omniglot
@@ -1193,6 +1243,9 @@ def _create_segmentation(self, size):
     def _file_stem(self, idx):
         return f"2008_{idx:06d}"
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self, config=dict(mode="segmentation"))
+
 
 class FakeDataTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.FakeData
@@ -1434,6 +1487,9 @@ def _magic(self, dtype, dims):
     def _encode(self, v):
         return torch.tensor(v, dtype=torch.int32).numpy().tobytes()[::-1]
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class FashionMNISTTestCase(MNISTTestCase):
     DATASET_CLASS = datasets.FashionMNIST
@@ -1585,6 +1641,9 @@ def test_classes(self, config):
             assert len(dataset.classes) == len(info["classes"])
             assert all([a == b for a, b in zip(dataset.classes, info["classes"])])
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class ImageFolderTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.ImageFolder
@@ -1606,6 +1665,9 @@ def test_classes(self, config):
             assert len(dataset.classes) == len(info["classes"])
             assert all([a == b for a, b in zip(dataset.classes, info["classes"])])
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class KittiTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.Kitti
@@ -1642,6 +1704,9 @@ def inject_fake_data(self, tmpdir, config):
 
         return split_to_num_examples[config["train"]]
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self, supports_target_keys=True)
+
 
 class SvhnTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.SVHN
@@ -2516,6 +2581,9 @@ def _meta_to_split_and_classification_ann(self, meta, idx):
         breed_id = "-1"
         return (image_id, class_id, species, breed_id)
 
+    def test_transforms_v2_wrapper(self):
+        datasets_utils.check_transforms_v2_wrapper(self)
+
 
 class StanfordCarsTestCase(datasets_utils.ImageDatasetTestCase):
     DATASET_CLASS = datasets.StanfordCars