add automatic feature type dispatch to functional transforms (#5323)

* add auto dispatch

* fix missing arguments error message

* remove pil kernel for erase

* automate feature specific parameter detection

* fix typos

* cleanup dispatcher call

* remove __torch_function__ from transform dispatch

* remove auto-generation

* revert unrelated changes

* remove implements decorator

* change register parameter order

* change order of transforms for readability

* add documentation for __torch_function__

* fix mypy

* inline check for support

* refactor kernel registering process

* refactor dispatch to be a regular decorator

* split kernels and dispatchers

* remove sentinels

* replace pass with ...

* appease mypy

* make single kernel dispatchers more concise

* make dispatcher signatures more generic

* make kernel checking more strict

* revert doc changes

* address Franciscos comments

* remove inplace

* rename kernel test module

* fix inplace

* remove special casing for pil and vanilla tensors

* address comments

* update docs

Author: pmeier

test/test_prototype_transforms_functional.py renamed to test/test_prototype_transforms_kernels.py

@@ -3,7 +3,7 @@
 
 import pytest
 import torch.testing
-import torchvision.prototype.transforms.functional as F
+import torchvision.prototype.transforms.kernels as K
 from torch import jit
 from torchvision.prototype import features
 
@@ -115,7 +115,7 @@ def __init__(self, *args, **kwargs):
 class KernelInfo:
     def __init__(self, name, *, sample_inputs_fn):
         self.name = name
-        self.kernel = getattr(F, name)
+        self.kernel = getattr(K, name)
         self._sample_inputs_fn = sample_inputs_fn
 
     def sample_inputs(self):
@@ -146,16 +146,16 @@ def horizontal_flip_image():
 @register_kernel_info_from_sample_inputs_fn
 def horizontal_flip_bounding_box():
     for bounding_box in make_bounding_boxes(formats=[features.BoundingBoxFormat.XYXY]):
-        yield SampleInput(bounding_box, image_size=bounding_box.image_size)
+        yield SampleInput(bounding_box, format=bounding_box.format, image_size=bounding_box.image_size)
 
 
 @register_kernel_info_from_sample_inputs_fn
 def resize_image():
     for image, interpolation in itertools.product(
         make_images(),
         [
-            F.InterpolationMode.BILINEAR,
-            F.InterpolationMode.NEAREST,
+            K.InterpolationMode.BILINEAR,
+            K.InterpolationMode.NEAREST,
         ],
     ):
         height, width = image.shape[-2:]

torchvision/prototype/features/_bounding_box.py

@@ -40,7 +40,7 @@ def __new__(
 
     def to_format(self, format: Union[str, BoundingBoxFormat]) -> BoundingBox:
         # import at runtime to avoid cyclic imports
-        from torchvision.prototype.transforms.functional import convert_bounding_box_format
+        from torchvision.prototype.transforms.kernels import convert_bounding_box_format
 
         if isinstance(format, str):
             format = BoundingBoxFormat[format]

torchvision/prototype/features/_encoded.py

@@ -40,7 +40,7 @@ def image_size(self) -> Tuple[int, int]:
 
     def decode(self) -> Image:
         # import at runtime to avoid cyclic imports
-        from torchvision.prototype.transforms.functional import decode_image_with_pil
+        from torchvision.prototype.transforms.kernels import decode_image_with_pil
 
         return Image(decode_image_with_pil(self))
 

torchvision/prototype/features/_feature.py

@@ -1,7 +1,7 @@
-from typing import Any, cast, Dict, Set, TypeVar, Union, Optional, Type, Callable
+from typing import Any, cast, Dict, Set, TypeVar, Union, Optional, Type, Callable, Tuple, Sequence, Mapping
 
 import torch
-from torch._C import _TensorBase
+from torch._C import _TensorBase, DisableTorchFunction
 
 
 F = TypeVar("F", bound="Feature")
@@ -76,5 +76,45 @@ def new_like(
         _metadata.update(metadata)
         return cls(data, dtype=dtype or other.dtype, device=device or other.device, **_metadata)
 
+    @classmethod
+    def __torch_function__(
+        cls,
+        func: Callable[..., torch.Tensor],
+        types: Tuple[Type[torch.Tensor], ...],
+        args: Sequence[Any] = (),
+        kwargs: Optional[Mapping[str, Any]] = None,
+    ) -> torch.Tensor:
+        """For general information about how the __torch_function__ protocol works,
+        see https://pytorch.org/docs/stable/notes/extending.html#extending-torch
+
+        TL;DR: Every time a PyTorch operator is called, it goes through the inputs and looks for the
+        ``__torch_function__`` method. If one is found, it is invoked with the operator as ``func`` as well as the
+        ``args`` and ``kwargs`` of the original call.
+
+        The default behavior of :class:`~torch.Tensor`'s is to retain a custom tensor type. For the :class:`Feature`
+        use case, this has two downsides:
+
+        1. Since some :class:`Feature`'s require metadata to be constructed, the default wrapping, i.e.
+           ``return cls(func(*args, **kwargs))``, will fail for them.
+        2. For most operations, there is no way of knowing if the input type is still valid for the output.
+
+        For these reasons, the automatic output wrapping is turned off for most operators.
+
+        Exceptions to this are:
+
+        - :func:`torch.clone`
+        - :meth:`torch.Tensor.to`
+        """
+        kwargs = kwargs or dict()
+        with DisableTorchFunction():
+            output = func(*args, **kwargs)
+
+        if func is torch.Tensor.clone:
+            return cls.new_like(args[0], output)
+        elif func is torch.Tensor.to:
+            return cls.new_like(args[0], output, dtype=output.dtype, device=output.device)
+        else:
+            return output
+
     def __repr__(self) -> str:
         return cast(str, torch.Tensor.__repr__(self)).replace("tensor", type(self).__name__)

torchvision/prototype/transforms/__init__.py

@@ -1,4 +1,5 @@
-from . import functional
-from .functional import InterpolationMode  # usort: skip
+from . import kernels  # usort: skip
+from . import functional  # usort: skip
+from .kernels import InterpolationMode  # usort: skip
 
 from ._presets import CocoEval, ImageNetEval, VocEval, Kinect400Eval, RaftEval

torchvision/prototype/transforms/functional/__init__.py

@@ -1,27 +1,14 @@
-from ._augment import erase_image, mixup_image, mixup_one_hot_label, cutmix_image, cutmix_one_hot_label
+from ._augment import erase, mixup, cutmix
 from ._color import (
-    adjust_brightness_image,
-    adjust_contrast_image,
-    adjust_saturation_image,
-    adjust_sharpness_image,
-    posterize_image,
-    solarize_image,
-    autocontrast_image,
-    equalize_image,
-    invert_image,
+    adjust_brightness,
+    adjust_contrast,
+    adjust_saturation,
+    adjust_sharpness,
+    posterize,
+    solarize,
+    autocontrast,
+    equalize,
+    invert,
 )
-from ._geometry import (
-    horizontal_flip_bounding_box,
-    horizontal_flip_image,
-    resize_bounding_box,
-    resize_image,
-    resize_segmentation_mask,
-    center_crop_image,
-    resized_crop_image,
-    InterpolationMode,
-    affine_image,
-    rotate_image,
-)
-from ._meta_conversion import convert_color_space, convert_bounding_box_format
-from ._misc import normalize_image
-from ._type_conversion import decode_image_with_pil, decode_video_with_av, label_to_one_hot
+from ._geometry import horizontal_flip, resize, center_crop, resized_crop, affine, rotate
+from ._misc import normalize

torchvision/prototype/transforms/functional/_augment.py

@@ -1,52 +1,57 @@
-from typing import Tuple
+from typing import TypeVar, Any
 
 import torch
+from torchvision.prototype import features
+from torchvision.prototype.transforms import kernels as K
 from torchvision.transforms import functional as _F
 
+from ._utils import dispatch
 
-erase_image = _F.erase
+T = TypeVar("T", bound=features.Feature)
 
 
-def _mixup(input: torch.Tensor, batch_dim: int, lam: float, inplace: bool) -> torch.Tensor:
-    if not inplace:
-        input = input.clone()
+@dispatch(
+    {
+        torch.Tensor: _F.erase,
+        features.Image: K.erase_image,
+    }
+)
+def erase(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
 
-    input_rolled = input.roll(1, batch_dim)
-    return input.mul_(lam).add_(input_rolled.mul_(1 - lam))
 
+@dispatch(
+    {
+        features.Image: K.mixup_image,
+        features.OneHotLabel: K.mixup_one_hot_label,
+    }
+)
+def mixup(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
 
-def mixup_image(image_batch: torch.Tensor, *, lam: float, inplace: bool = False) -> torch.Tensor:
-    if image_batch.ndim < 4:
-        raise ValueError("Need a batch of images")
 
-    return _mixup(image_batch, -4, lam, inplace)
+@dispatch(
+    {
+        features.Image: K.cutmix_image,
+        features.OneHotLabel: K.cutmix_one_hot_label,
+    }
+)
+def cutmix(input: T, *args: Any, **kwargs: Any) -> T:
+    """Perform the CutMix operation as introduced in the paper
+    `"CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features" <https://arxiv.org/abs/1905.04899>`_.
 
+    Dispatch to the corresponding kernels happens according to this table:
 
-def mixup_one_hot_label(one_hot_label_batch: torch.Tensor, *, lam: float, inplace: bool = False) -> torch.Tensor:
-    if one_hot_label_batch.ndim < 2:
-        raise ValueError("Need a batch of one hot labels")
+    .. table::
+       :widths: 30 70
 
-    return _mixup(one_hot_label_batch, -2, lam, inplace)
+       ====================================================  ================================================================
+       :class:`~torchvision.prototype.features.Image`        :func:`~torch.prototype.transforms.kernels.cutmix_image`
+       :class:`~torchvision.prototype.features.OneHotLabel`  :func:`~torch.prototype.transforms.kernels.cutmix_one_hot_label`
+       ====================================================  ================================================================
 
-
-def cutmix_image(image_batch: torch.Tensor, *, box: Tuple[int, int, int, int], inplace: bool = False) -> torch.Tensor:
-    if image_batch.ndim < 4:
-        raise ValueError("Need a batch of images")
-
-    if not inplace:
-        image_batch = image_batch.clone()
-
-    x1, y1, x2, y2 = box
-    image_rolled = image_batch.roll(1, -4)
-
-    image_batch[..., y1:y2, x1:x2] = image_rolled[..., y1:y2, x1:x2]
-    return image_batch
-
-
-def cutmix_one_hot_label(
-    one_hot_label_batch: torch.Tensor, *, lam_adjusted: float, inplace: bool = False
-) -> torch.Tensor:
-    if one_hot_label_batch.ndim < 2:
-        raise ValueError("Need a batch of one hot labels")
-
-    return _mixup(one_hot_label_batch, -2, lam_adjusted, inplace)
+    Please refer to the kernel documentations for a detailed explanation of the functionality and parameters.
+    """
+    ...

torchvision/prototype/transforms/functional/_color.py

@@ -1,20 +1,119 @@
+from typing import TypeVar, Any
+
+import PIL.Image
+import torch
+from torchvision.prototype import features
+from torchvision.prototype.transforms import kernels as K
 from torchvision.transforms import functional as _F
 
+from ._utils import dispatch
+
+T = TypeVar("T", bound=features.Feature)
+
+
+@dispatch(
+    {
+        torch.Tensor: _F.adjust_brightness,
+        PIL.Image.Image: _F.adjust_brightness,
+        features.Image: K.adjust_brightness_image,
+    }
+)
+def adjust_brightness(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
+
+
+@dispatch(
+    {
+        torch.Tensor: _F.adjust_saturation,
+        PIL.Image.Image: _F.adjust_saturation,
+        features.Image: K.adjust_saturation_image,
+    }
+)
+def adjust_saturation(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
+
+
+@dispatch(
+    {
+        torch.Tensor: _F.adjust_contrast,
+        PIL.Image.Image: _F.adjust_contrast,
+        features.Image: K.adjust_contrast_image,
+    }
+)
+def adjust_contrast(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
+
+
+@dispatch(
+    {
+        torch.Tensor: _F.adjust_sharpness,
+        PIL.Image.Image: _F.adjust_sharpness,
+        features.Image: K.adjust_sharpness_image,
+    }
+)
+def adjust_sharpness(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
+
 
-adjust_brightness_image = _F.adjust_brightness
+@dispatch(
+    {
+        torch.Tensor: _F.posterize,
+        PIL.Image.Image: _F.posterize,
+        features.Image: K.posterize_image,
+    }
+)
+def posterize(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
 
-adjust_saturation_image = _F.adjust_saturation
 
-adjust_contrast_image = _F.adjust_contrast
+@dispatch(
+    {
+        torch.Tensor: _F.solarize,
+        PIL.Image.Image: _F.solarize,
+        features.Image: K.solarize_image,
+    }
+)
+def solarize(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
 
-adjust_sharpness_image = _F.adjust_sharpness
 
-posterize_image = _F.posterize
+@dispatch(
+    {
+        torch.Tensor: _F.autocontrast,
+        PIL.Image.Image: _F.autocontrast,
+        features.Image: K.autocontrast_image,
+    }
+)
+def autocontrast(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
 
-solarize_image = _F.solarize
 
-autocontrast_image = _F.autocontrast
+@dispatch(
+    {
+        torch.Tensor: _F.equalize,
+        PIL.Image.Image: _F.equalize,
+        features.Image: K.equalize_image,
+    }
+)
+def equalize(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...
 
-equalize_image = _F.equalize
 
-invert_image = _F.invert
+@dispatch(
+    {
+        torch.Tensor: _F.invert,
+        PIL.Image.Image: _F.invert,
+        features.Image: K.invert_image,
+    }
+)
+def invert(input: T, *args: Any, **kwargs: Any) -> T:
+    """ADDME"""
+    ...