only return small set of targets by default from dataset wrapper

[pmeier]

From benchmarking, it seems that the detection references with v2 are quite a bit slower than v1 for two reasons:

1. The input sample for COCO is fairly large with roughly > 50 items. The largest part come from the "segmentations" key, which stores the vertex coordinates of polygons in nested lists. Recursing through them every time with pytree is slow.
2. The "segmentations" key will be decoded into a datapoints.Mask and later on transformed although for regular object detection, this is not needed.

This PR introduces the target_keys parameter that let's users select which keys they actually want.

ToDo:

• add support for "all" if users want everything
• document this
• honor target_keys in all irregular datasets besides CocoDetection

cc @vfdev-5

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/7488

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 Failures

As of commit 47f638d:

NEW FAILURES - The following jobs have failed:

• pytorch/vision / conda-py3_8-cpu (gh)
• pytorch/vision / manywheel-py3_9-rocm5_3 (gh)
• pytorch/vision / wheel-py3_10-cpu (gh)
• pytorch/vision / wheel-py3_9-cpu (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[pmeier]

We need the full sample now for target_keys="all".

[NicolasHug]

Thanks Philip, some comments below but looks good

[NicolasHug]

Thanks Philip, LGTM

[pmeier]

Sneak peak into the performance gains of this change

Summaries

           v2 / v1
Tensor        0.32
PIL           0.33

                     [a]   [b]   [c]   [d]   [e]
   Tensor, v1, [a]  1.00  3.16  0.97  2.91  2.79
   Tensor, v2, [b]  0.32  1.00  0.31  0.92  0.88
      PIL, v1, [c]  1.03  3.26  1.00  2.99  2.87
      PIL, v2, [d]  0.34  1.09  0.33  1.00  0.96
Datapoint, v2, [e]  0.36  1.14  0.35  1.04  1.00

Slowdown as row / col

Details

############################################################
detection-ssdlite
############################################################
input_type='Tensor', api_version='v1'

Results computed for 1_000 samples

                          median          std   
ConvertCocoPolysToMaskV1    2663 µs +-   4269 µs
PILToTensorV1                251 µs +-     71 µs
RandomIoUCropV1              456 µs +-   6950 µs
RandomHorizontalFlipV1        17 µs +-    200 µs
ConvertImageDtypeV1          271 µs +-    172 µs

total                       3657 µs
------------------------------------------------------------
input_type='Tensor', api_version='v2'

Results computed for 1_000 samples

                               median          std   
WrapCocoSampleForTransformsV2      61 µs +-     34 µs
PILToTensor                       218 µs +-     46 µs
RandomIoUCrop                     394 µs +-   6302 µs
RandomHorizontalFlip              132 µs +-    116 µs
ConvertDtype                      178 µs +-    128 µs
SanitizeBoundingBox               173 µs +-     14 µs

total                            1156 µs
------------------------------------------------------------
input_type='PIL', api_version='v1'

Results computed for 1_000 samples

                          median          std   
ConvertCocoPolysToMaskV1    2735 µs +-   4329 µs
RandomIoUCropV1              548 µs +-   7064 µs
RandomHorizontalFlipV1        19 µs +-    223 µs
PILToTensorV1                182 µs +-    118 µs
ConvertImageDtypeV1          282 µs +-    178 µs

total                       3767 µs
------------------------------------------------------------
input_type='PIL', api_version='v2'

Results computed for 1_000 samples

                               median          std   
WrapCocoSampleForTransformsV2      64 µs +-     26 µs
RandomIoUCrop                     489 µs +-   6620 µs
RandomHorizontalFlip              143 µs +-    126 µs
PILToTensor                       189 µs +-     84 µs
ConvertDtype                      191 µs +-    121 µs
SanitizeBoundingBox               182 µs +-     16 µs

total                            1259 µs
------------------------------------------------------------
input_type='Datapoint', api_version='v2'

Results computed for 1_000 samples

                               median          std   
WrapCocoSampleForTransformsV2      65 µs +-     21 µs
ToImageTensor                     272 µs +-     92 µs
RandomIoUCrop                     438 µs +-   6962 µs
RandomHorizontalFlip              148 µs +-    183 µs
ConvertDtype                      205 µs +-    162 µs
SanitizeBoundingBox               185 µs +-     25 µs

total                            1313 µs

Roughly 3x over v1. To be fair, v1 is affected by #7489 so this comparison is not totally fair. However, it accurately reflects the current state of our references vs. what can be achieved with v2.