[fbsync] only return small set of targets by default from dataset wrapper (#7488)

Reviewed By: vmoens

Differential Revision: D45183665

fbshipit-source-id: 1d803dee2f2e1442113ad8d0e7f95e3b1314f7be

Author: NicolasHug

gallery/plot_transforms_v2_e2e.py

@@ -75,15 +75,16 @@ def load_example_coco_detection_dataset(**kwargs):
 # :func:`~torchvision.datasets.wrap_dataset_for_transforms_v2` function. For
 # :class:`~torchvision.datasets.CocoDetection`, this changes the target structure to a single dictionary of lists. It
 # also adds the key-value-pairs ``"boxes"``, ``"masks"``, and ``"labels"`` wrapped in the corresponding
-# ``torchvision.datapoints``.
+# ``torchvision.datapoints``. By default, it only returns ``"boxes"`` and ``"labels"`` to avoid transforming unnecessary
+# items down the line, but you can pass the ``target_type`` parameter for fine-grained control.
 
 dataset = datasets.wrap_dataset_for_transforms_v2(dataset)
 
 sample = dataset[0]
 image, target = sample
 print(type(image))
 print(type(target), list(target.keys()))
-print(type(target["boxes"]), type(target["masks"]), type(target["labels"]))
+print(type(target["boxes"]), type(target["labels"]))
 
 ########################################################################################################################
 # As baseline, let's have a look at a sample without transformations:

test/datasets_utils.py

@@ -572,9 +572,21 @@ def test_transforms_v2_wrapper(self, config):
 
         try:
             with self.create_dataset(config) as (dataset, _):
-                wrapped_dataset = wrap_dataset_for_transforms_v2(dataset)
-                wrapped_sample = wrapped_dataset[0]
-                assert tree_any(lambda item: isinstance(item, (Datapoint, PIL.Image.Image)), wrapped_sample)
+                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)
+                    wrapped_sample = wrapped_dataset[0]
+
+                    assert tree_any(lambda item: isinstance(item, (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):

test/test_datasets.py

@@ -771,6 +771,8 @@ def _create_annotations(self, image_ids, num_annotations_per_image):
                     bbox=torch.rand(4).tolist(),
                     segmentation=[torch.rand(8).tolist()],
                     category_id=int(torch.randint(91, ())),
+                    area=float(torch.rand(1)),
+                    iscrowd=int(torch.randint(2, size=(1,))),
                 )
             )
             annotion_id += 1
@@ -3336,7 +3338,7 @@ def test_subclass(self, mocker):
         mocker.patch.dict(
             datapoints._dataset_wrapper.WRAPPER_FACTORIES,
             clear=False,
-            values={datasets.FakeData: lambda dataset: lambda idx, sample: sentinel},
+            values={datasets.FakeData: lambda dataset, target_keys: lambda idx, sample: sentinel},
         )
 
         class MyFakeData(datasets.FakeData):