Repository: pytorch/vision

Issue:
The prototype `FiveCrop`/`TenCrop`/`BatchMultiCrop` produce unexpected results
### 🐛 Describe the bug

The current implementation of the transforms produce unexpected results for Classification.

- If a single image comes in (no batch&collated), a batched result comes out with a single label
- if a batch of images comes in, we get a new image batch with 5x the original size and a label vector with length equal to the original size

In both cases, the `X` and the `Y` lengths don't match. This creates issues on validation pipelines. 

One approach would be to duplicate the length of Labels to ensure they match. It's quite likely the same approach should be considered for meta-data (like ids) and other content included in the record. Things get even more complex if we consider extending the transform to Detection as we would need to crop the Masks/BBoxes, Cleanup the BBoxes and sync the Labels for each of the new crops.

This complexity is possible the reason why existing the `FiveCrop` implementation on stable don't offer a way to stack the result for the user but instead it just provides an [example](https://github.com/pytorch/vision/blob/main/torchvision/transforms/transforms.py#L992-L1000) in the documentation on how it can be done. Due to the above, I believe that the current implementation of `BatchMultiCrop` is flawed and should be either removed or redeveloped.

### Versions

Latest main branch

cc @vfdev-5 @datumbox @bjuncek @pmeier


Discussion:
I'm ok with the removal of `BatchMultiCrop` for the reasons given above. In that case we need to have at least a minimal example in the documentation how batching can be done. I came up with this for a simple classification example with `features.Image`'s and `features.Label`'s:

```py
import torch
from torchvision.prototype import transforms, features

from torch.utils._pytree import tree_flatten


class BatchMultiCrop(transforms.Transform):
    _transformed_types = (transforms.MultiCropResult, features.Label)

    def _get_params(self, sample):
        sample_flat, _ = tree_flatten(sample)
        batch_size = len(
            next(item for item in sample_flat if isinstance(item, transforms.MultiCropResult))
        )
        return dict(batch_size=batch_size)

    def _transform(self, inpt, params):
        if isinstance(inpt, transforms.MultiCropResult):
            output = torch.stack(inpt)
            return features.Image.new_like(inpt[0], output)
        else:  # isinstance(inpt, features.Label):
            output = inpt.repeat(params["batch_size"], *[1] * inpt.ndim)
            return features.Label.new_like(inpt, output)

transform = transform.Compose([transforms.FiveCrop(), BatchMultiCrop()])
```

This is already quite long for an minimal example, but I don't think we can get it much shorter. There are a few options though:
- remove the `torch` and `torchvision` imports as they can be assumed to be present.
- increase the density, e.g. instead of having an `output` variable in `_transform()`, the computation could happen directly in the `new_like` kernel. Normally this hinders readability, but the computation is fairly simple here.
- Instead of dynamically extracting the `batch_size`, this could also be given as fixed parameter. That way we would need an `__init__`, but could save a few lines as well as an import while overall making the example easier.

Apart from that I have two more comments:
- If we keep the dynamic `batch_size` extraction, we need to rely on `tree_flatten`. I dislike showing off an example where the user needs to use "private" functionality from another package, namely `torch.utils._pytree`.
- The name `MultiCropResult` is somewhat misleading, since it only contains images rather than a more general "result". If we go through with removing `BatchMultiCrop`, we need to expose `MultiCropResult` so users can write their own batching transforms. In that case, we should probably rename it to avoid confusion. Maybe `MultiCropImages`?
I'm OK with the example you provide. I don't mind that the example is on the long side (I wouldn't try to make it more dense) but I do agree about the `tree_flatten` remark. Perhaps what we can do in this API is avoid writing a generic `BatchMultiCrop` class and instead provide a more specialised implementation for classification which expects specific input types on `forward(img: Image, lbl: Label)` and thus removes the extra complexity of packing/unpacking. This obviously make supporting meta-data harder but again the intention of the example is to show a simple use-case of how to handle the output. Thoughts?
My only concern with having `forward(img: Image, lbl: Label)` is that it won't work with our datasets whereas `FiveCrop` / `TenCrop` do out of the box.
Yes, it's not a generic solution I agree. But because currently we can't offer a generic solution, I think it's worth being as specific and clear as possible on our example. On the future we could revise this and try to offer a truly generic class that handles all tasks.

@vfdev-5 Thoughts?
The other option to avoid the `tree_flatten` call is to take the `batch_size` as parameter:

```py
class BatchMultiCrop(transforms.Transform):
    _transformed_types = (transforms.MultiCropResult, features.Label)

    def __init__(self, batch_size):
        super().__init__()
        self.batch_size = batch_size

    def _transform(self, inpt, params):
        if isinstance(inpt, transforms.MultiCropResult):
            output = torch.stack(inpt)
            return features.Image.new_like(inpt[0], output)
        else:  # isinstance(inpt, features.Label):
            output = inpt.repeat(self.batch_size, *[1] * inpt.ndim)
            return features.Label.new_like(inpt, output)

transform = transform.Compose([transforms.FiveCrop(), BatchMultiCrop(5)])
```

That way arbitrary input structures are allowed and we don't need to showcase an example where our API is not sufficient and one needs to overwrite `forward`.
Passing the batch-size is very weird as it's not something we do anywhere else. IMO given that we don't offer a proper implementation of `BatchMultiCrop` out of the box, we should aim for simplicity. This means show to the user the simplest most specific way to handle the use-case and let them build on top of it.
That would look like

```py
class BatchMultiCrop(transforms.Transform):
    def forward(self, sample: Tuple[transforms.MultiCropResult, features.Label]):
        images, labels = sample
        batch_size = len(images)
        images = features.Image.new_like(images[0], torch.stack(images))
        labels = features.Label.new_like(labels, labels.repeat(batch_size, *[1] * labels.ndim))
        return images, labels
```

Note that we cannot have a `forward(self, images: transforms.MultiCropResult, labels: features.Label)` signature, because even if you use `transform = transforms.Compose([transforms.FiveCrop(), BatchMultiCrop()])` with `transform(image, label)`, the output of `FiveCrop` will always be a tuple. The only way to avoid this would be to remove the `FiveCrop` from the example, but I'm not sure if we aren't then wading into "useless" example since we showcase something that will not happen in practice and move the burden entirely on the user.
Looks OK to me.

Another thing I would avoid, if you agree, is the use of `transforms.MultiCropResult`. Why can't we put everything a simple List that clearly shows to the signature what is the expected input (`List[features.Image`). I understand that on the future, if we want to provide such a transform, we would need to be able to identify that this is a special type of list. But if this list is not used at the moment, it only gets in the way for the user. Such a future change can happen in a BC manner I think.

Let me know if I'm missing something important here.

Resolving patch:
diff --git a/torchvision/prototype/transforms/__init__.py b/torchvision/prototype/transforms/__init__.py
--- a/torchvision/prototype/transforms/__init__.py
+++ b/torchvision/prototype/transforms/__init__.py
@@ -16,7 +16,6 @@
 )
 from ._container import Compose, RandomApply, RandomChoice, RandomOrder
 from ._geometry import (
-    BatchMultiCrop,
     CenterCrop,
     ElasticTransform,
     FiveCrop,
diff --git a/torchvision/prototype/transforms/_geometry.py b/torchvision/prototype/transforms/_geometry.py
--- a/torchvision/prototype/transforms/_geometry.py
+++ b/torchvision/prototype/transforms/_geometry.py
@@ -8,7 +8,7 @@
 from torchvision.ops.boxes import box_iou
 from torchvision.prototype import features
 from torchvision.prototype.transforms import functional as F, Transform
-from torchvision.transforms.functional import InterpolationMode, pil_to_tensor
+from torchvision.transforms.functional import InterpolationMode
 from torchvision.transforms.functional_tensor import _parse_pad_padding
 from torchvision.transforms.transforms import _check_sequence_input, _setup_angle, _setup_size
 
@@ -136,30 +136,41 @@ def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
         )
 
 
-class MultiCropResult(list):
-    """Helper class for :class:`~torchvision.prototype.transforms.BatchMultiCrop`.
-
-    Outputs of multi crop transforms such as :class:`~torchvision.prototype.transforms.FiveCrop` and
-    `:class:`~torchvision.prototype.transforms.TenCrop` should be wrapped in this in order to be batched correctly by
-    :class:`~torchvision.prototype.transforms.BatchMultiCrop`.
+class FiveCrop(Transform):
+    """
+    Example:
+        >>> class BatchMultiCrop(transforms.Transform):
+        ...     def forward(self, sample: Tuple[Tuple[features.Image, ...], features.Label]):
+        ...         images, labels = sample
+        ...         batch_size = len(images)
+        ...         images = features.Image.new_like(images[0], torch.stack(images))
+        ...         labels = features.Label.new_like(labels, labels.repeat(batch_size))
+        ...         return images, labels
+        ...
+        >>> image = features.Image(torch.rand(3, 256, 256))
+        >>> label = features.Label(0)
+        >>> transform = transforms.Compose([transforms.FiveCrop(), BatchMultiCrop()])
+        >>> images, labels = transform(image, label)
+        >>> images.shape
+        torch.Size([5, 3, 224, 224])
+        >>> labels.shape
+        torch.Size([5])
     """
 
-    pass
-
-
-class FiveCrop(Transform):
     def __init__(self, size: Union[int, Sequence[int]]) -> None:
         super().__init__()
         self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
 
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
+        # TODO: returning a list is technically BC breaking since FiveCrop returned a tuple before. We switched to a
+        #  list here to align it with TenCrop.
         if isinstance(inpt, features.Image):
             output = F.five_crop_image_tensor(inpt, self.size)
-            return MultiCropResult(features.Image.new_like(inpt, o) for o in output)
+            return tuple(features.Image.new_like(inpt, o) for o in output)
         elif is_simple_tensor(inpt):
-            return MultiCropResult(F.five_crop_image_tensor(inpt, self.size))
+            return F.five_crop_image_tensor(inpt, self.size)
         elif isinstance(inpt, PIL.Image.Image):
-            return MultiCropResult(F.five_crop_image_pil(inpt, self.size))
+            return F.five_crop_image_pil(inpt, self.size)
         else:
             return inpt
 
@@ -171,6 +182,10 @@ def forward(self, *inputs: Any) -> Any:
 
 
 class TenCrop(Transform):
+    """
+    See :class:`~torchvision.prototype.transforms.FiveCrop` for an example.
+    """
+
     def __init__(self, size: Union[int, Sequence[int]], vertical_flip: bool = False) -> None:
         super().__init__()
         self.size = _setup_size(size, error_msg="Please provide only two dimensions (h, w) for size.")
@@ -179,11 +194,11 @@ def __init__(self, size: Union[int, Sequence[int]], vertical_flip: bool = False)
     def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
         if isinstance(inpt, features.Image):
             output = F.ten_crop_image_tensor(inpt, self.size, vertical_flip=self.vertical_flip)
-            return MultiCropResult(features.Image.new_like(inpt, o) for o in output)
+            return [features.Image.new_like(inpt, o) for o in output]
         elif is_simple_tensor(inpt):
-            return MultiCropResult(F.ten_crop_image_tensor(inpt, self.size, vertical_flip=self.vertical_flip))
+            return F.ten_crop_image_tensor(inpt, self.size, vertical_flip=self.vertical_flip)
         elif isinstance(inpt, PIL.Image.Image):
-            return MultiCropResult(F.ten_crop_image_pil(inpt, self.size, vertical_flip=self.vertical_flip))
+            return F.ten_crop_image_pil(inpt, self.size, vertical_flip=self.vertical_flip)
         else:
             return inpt
 
@@ -194,22 +209,6 @@ def forward(self, *inputs: Any) -> Any:
         return super().forward(sample)
 
 
-class BatchMultiCrop(Transform):
-    _transformed_types = (MultiCropResult,)
-
-    def _transform(self, inpt: Any, params: Dict[str, Any]) -> Any:
-        crops = inpt
-        if isinstance(inpt[0], PIL.Image.Image):
-            crops = [pil_to_tensor(crop) for crop in crops]
-
-        batch = torch.stack(crops)
-
-        if isinstance(inpt[0], features.Image):
-            batch = features.Image.new_like(inpt[0], batch)
-
-        return batch
-
-
 def _check_fill_arg(fill: Union[int, float, Sequence[int], Sequence[float]]) -> None:
     if not isinstance(fill, (numbers.Number, tuple, list)):
         raise TypeError("Got inappropriate fill arg")


Test patch: