reinstate old test

Author: pmeier

test/datasets_utils.py

@@ -171,46 +171,6 @@ 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
-
-    def check_wrapped_samples(dataset):
-        for wrapped_sample in dataset:
-            assert tree_any(
-                lambda item: isinstance(item, (datapoints.Image, datapoints.Video, PIL.Image.Image)), wrapped_sample
-            )
-
-    target_keyss = [None]
-    if supports_target_keys:
-        target_keyss.append("all")
-
-    for target_keys in target_keyss:
-        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"]
-
-            check_wrapped_samples(wrapped_dataset)
-
-    # On macOS, forking for multiprocessing is not available and thus spawning is used by default. For this to work,
-    # the whole pipeline including the dataset needs to be pickleable, which is what we are enforcing here.
-    if platform.system() == "Darwin":
-        with dataset_test_case.create_dataset(config) as (dataset, _):
-            wrapped_dataset = wrap_dataset_for_transforms_v2(dataset)
-            dataloader = DataLoader(
-                wrapped_dataset, num_workers=2, multiprocessing_context="spawn", collate_fn=_no_collate
-            )
-
-            check_wrapped_samples(dataloader)
-
-
 class DatasetTestCase(unittest.TestCase):
     """Abstract base class for all dataset testcases.
 
@@ -606,6 +566,42 @@ 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):
+                for target_keys in [None, "all"]:
+                    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 = 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.Datapoint, PIL.Image.Image)), wrapped_sample
+                    )
+        except TypeError as error:
+            msg = f"No wrapper exists for dataset class {type(dataset).__name__}"
+            if str(error).startswith(msg):
+                pytest.skip(msg)
+            raise error
+        except RuntimeError as error:
+            if "currently not supported by this wrapper" in str(error):
+                pytest.skip("Config is currently not supported by this wrapper")
+            raise error
+
 
 class ImageDatasetTestCase(DatasetTestCase):
     """Abstract base class for image dataset testcases.
@@ -687,6 +683,40 @@ 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 _no_collate(batch):
+    return batch
+
+
+def check_transforms_v2_wrapper_spawn(dataset):
+    # On Linux and Windows, the DataLoader forks the main process by default. This is not available on macOS, so new
+    # subprocesses are spawned. This requires the whole pipeline including the dataset to be pickleable, which is what
+    # we are enforcing here.
+    if platform.system() != "Darwin":
+        pytest.skip("Multiprocessing spawning is only checked on macOS.")
+
+    from torch.utils.data import DataLoader
+    from torchvision import datapoints
+    from torchvision.datasets import wrap_dataset_for_transforms_v2
+
+    wrapped_dataset = wrap_dataset_for_transforms_v2(dataset)
+
+    dataloader = DataLoader(wrapped_dataset, num_workers=2, multiprocessing_context="spawn", 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
+        )
+
 
 def create_image_or_video_tensor(size: Sequence[int]) -> torch.Tensor:
     r"""Create a random uint8 tensor.

test/test_datasets.py

@@ -183,8 +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"))
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset(target_type="category") as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class Caltech256TestCase(datasets_utils.ImageDatasetTestCase):
@@ -261,8 +262,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)
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset() as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class CityScapesTestCase(datasets_utils.ImageDatasetTestCase):
@@ -388,9 +390,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):
+    def test_transforms_v2_wrapper_spawn(self):
         for target_type in ["instance", "semantic", ["instance", "semantic"]]:
-            datasets_utils.check_transforms_v2_wrapper(self, config=dict(target_type=target_type))
+            with self.create_dataset(target_type=target_type) as (dataset, _):
+                datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class ImageNetTestCase(datasets_utils.ImageDatasetTestCase):
@@ -423,8 +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)
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset() as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class CIFAR10TestCase(datasets_utils.ImageDatasetTestCase):
@@ -620,9 +624,10 @@ def test_images_names_split(self):
 
         assert merged_imgs_names == all_imgs_names
 
-    def test_transforms_v2_wrapper(self):
+    def test_transforms_v2_wrapper_spawn(self):
         for target_type in ["identity", "bbox", ["identity", "bbox"]]:
-            datasets_utils.check_transforms_v2_wrapper(self, config=dict(target_type=target_type))
+            with self.create_dataset(target_type=target_type) as (dataset, _):
+                datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class VOCSegmentationTestCase(datasets_utils.ImageDatasetTestCase):
@@ -711,8 +716,9 @@ def add_bndbox(obj, bndbox=None):
 
         return data
 
-    def test_transforms_v2_wrapper(self):
-        datasets_utils.check_transforms_v2_wrapper(self)
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset() as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class VOCDetectionTestCase(VOCSegmentationTestCase):
@@ -734,9 +740,9 @@ 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)
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset() as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class CocoDetectionTestCase(datasets_utils.ImageDatasetTestCase):
@@ -808,8 +814,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)
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset() as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class CocoCaptionsTestCase(CocoDetectionTestCase):
@@ -827,7 +834,7 @@ def test_captions(self):
             _, captions = dataset[0]
             assert tuple(captions) == tuple(info["captions"])
 
-    def test_transforms_v2_wrapper(self):
+    def test_transforms_v2_wrapper_spawn(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.")
@@ -998,8 +1005,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"))
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset(output_format="TCHW") as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class HMDB51TestCase(datasets_utils.VideoDatasetTestCase):
@@ -1228,8 +1236,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"))
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset(mode="segmentation") as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class FakeDataTestCase(datasets_utils.ImageDatasetTestCase):
@@ -1680,8 +1689,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)
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset() as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class SvhnTestCase(datasets_utils.ImageDatasetTestCase):
@@ -2557,8 +2567,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)
+    def test_transforms_v2_wrapper_spawn(self):
+        with self.create_dataset() as (dataset, _):
+            datasets_utils.check_transforms_v2_wrapper_spawn(dataset)
 
 
 class StanfordCarsTestCase(datasets_utils.ImageDatasetTestCase):