port FiveCrop and TenCrop to prototype API

[pmeier]

Both transforms are outliers in the current stable API as discussed in #5500 (comment). Even if we want to deprecate them in favor of a different approach eventually, we need to support them at day 0 of the roll out.

The issues with these transforms arise, because they always need to be followed by another to batch them together:

vision/torchvision/transforms/transforms.py

Lines 980 to 984 in 71d2bb0

	Example:
	>>> transform = Compose([
	>>> FiveCrop(size), # this is a list of PIL Images
	>>> Lambda(lambda crops: torch.stack([ToTensor()(crop) for crop in crops])) # returns a 4D tensor
	>>> ])

We cannot re-use this lambda approach since crops might be buried inside a nested container. Thus, we need a custom transformation that picks up on the result of such a multi crop transform and batches it. To communicate which elements need to be batched between the two transformed, I've opted to wrap the return values of the multi crop kernels in named tuples, because they are full JIT scriptable. The old FiveCrop returned a regular tuple, so this is 100% compatible. TenCrop returned a list, but this is still fine unless someone is using list specific methods like append or the like.

Example usage:

from torchvision.prototype import datasets, transforms, features

transform = transforms.Compose(
    transforms.DecodeImage(),
    transforms.TenCrop(16),
    transforms.BatchMultiCrop(),
)

for sample in datasets.load("imagenet").map(transform):
    print(sample["image"].shape)
    break

torch.Size([10, 3, 16, 16])

To test with PIL images, insert transforms.Lambda(to_pil_image, features.Image) after the decoding.

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit f1a5003 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[pmeier]

We need this addition to be able to exclude named tuples as sequences in the BatchMultiCrop transform. My gut says we are going to need this fine grained control again for other transforms that are not yet ported / implemented. If that turns out not to be true, I'm happy to simply implement a custom solution only in the BatchMultiCrop transform given that we probably deprecate it anyway.

[datumbox]

We should start from the assumption that this is not needed and add it later. Starting with a simple solution first is a good prior. Please simplify as much as possible.

[pmeier]

We need this exclude here, because named tuples by default would be recognized as sequence and thus we would only get the individual elements rather than everything at once.

[datumbox]

I think that's one more reason not to use named tuples.

[datumbox]

@pmeier As we discussed previously offline, the specific transforms are a bit problematic as they don't fit nicely on the new API.
The only reason why we provide support for them is for BC, so I would avoid introducing any new features. Let's offer the simplest possible solution just to ensure we won't break anyone's code.

[datumbox]

We should start from the assumption that this is not needed and add it later. Starting with a simple solution first is a good prior. Please simplify as much as possible.

[datumbox]

I think that's one more reason not to use named tuples.

[pmeier]

After some offline discussion @datumbox and I agreed that we indeed need the BatchMultiCrop transform. Nevertheless, with the new commits the implementation is a lot simply, by moving the "result wrapper" onto the transformations and being able to use only a single subtype of list rather than two namedtuples that would have been necessary to satisfy torchscript. Plus, instead of making apply_recursively more general, BatchMultiCrop now implements a custom version of it to keep it simple.

[datumbox]

LGTM.

I think we should also review on the near future whether the very detailed nested sample is something we need to support or if that's something we don't need on the Datasets. If it turns out it's not needed it will allow us to massively simplify the transforms.

cc @NicolasHug