Repository: pytorch/vision

Issue:
Support all integer and floating point dtypes in prototype transform kernels?
The standard rule for dtype support for images and videos is: 

- All floating point and integer tensors are supported.
- Floating point tensors are valid in the range `[0.0, 1.0]` and integer tensors in `[0, torch.iinfo(dtype).max]` (this is currently under review since there were a few cases, where this was not true or simply not handled. See #6825)

However we have currently two kernels that only support `uint8` images or videos:

- https://github.com/pytorch/vision/blob/c84dbfad97251271a789b252a2a1a52c73f623ff/torchvision/prototype/transforms/functional/_color.py#L373-L375
- https://github.com/pytorch/vision/blob/c84dbfad97251271a789b252a2a1a52c73f623ff/torchvision/transforms/functional_tensor.py#L788-L789

This also holds for transforms v1 so this is not a problem of the new API. 

One consequence of that is that AA transforms are only supported for `uint8` images

https://github.com/pytorch/vision/blob/c84dbfad97251271a789b252a2a1a52c73f623ff/torchvision/transforms/autoaugment.py#L104-L107

since both 

https://github.com/pytorch/vision/blob/c84dbfad97251271a789b252a2a1a52c73f623ff/torchvision/transforms/autoaugment.py#L76-L77

and 

https://github.com/pytorch/vision/blob/c84dbfad97251271a789b252a2a1a52c73f623ff/torchvision/transforms/autoaugment.py#L82-L83

are used.

One possible way of mitigating this to simply have a `convert_dtype(image, torch.uint8)` in the beginning and converting back after computation. 

That is probably needed for `equalize` since we recently switched away from the histogram ops of `torch` towards our "custom" implementation to enable batch processing (#6757). However, this relies on the fact that the input is an integer and in its current form even on `uint8` due to some hardcoded constants.

For `posterize` I think it is fairly easy to provide the same functionality for float inputs directly without going through a dtype conversion first.

cc @vfdev-5 @datumbox @bjuncek


Discussion:
I'm in favour of doing this. It will allow users to cast to float32 first which might be beneficial for cases where the uint8 kernels are slower.
Before closing this ticket, we should make the following proposed optimization on AugMix:
https://github.com/pytorch/vision/blob/e1f464bded7781be85fdf1762c01a7869d3aec97/torchvision/prototype/transforms/_auto_augment.py#L509-L513

Edit: I checked the above and doesn't actually improve the speed, so we should just remove the comment.

Resolving patch:
diff --git a/torchvision/prototype/transforms/_auto_augment.py b/torchvision/prototype/transforms/_auto_augment.py
--- a/torchvision/prototype/transforms/_auto_augment.py
+++ b/torchvision/prototype/transforms/_auto_augment.py
@@ -137,7 +137,8 @@ def _apply_image_or_video_transform(
         elif transform_id == "Posterize":
             return F.posterize(image, bits=int(magnitude))
         elif transform_id == "Solarize":
-            return F.solarize(image, threshold=magnitude)
+            bound = 1.0 if isinstance(image, torch.Tensor) and image.is_floating_point() else 255.0
+            return F.solarize(image, threshold=bound * magnitude)
         elif transform_id == "AutoContrast":
             return F.autocontrast(image)
         elif transform_id == "Equalize":
@@ -169,7 +170,7 @@ class AutoAugment(_AutoAugmentBase):
             lambda num_bins, height, width: (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4))).round().int(),
             False,
         ),
-        "Solarize": (lambda num_bins, height, width: torch.linspace(255.0, 0.0, num_bins), False),
+        "Solarize": (lambda num_bins, height, width: torch.linspace(1.0, 0.0, num_bins), False),
         "AutoContrast": (lambda num_bins, height, width: None, False),
         "Equalize": (lambda num_bins, height, width: None, False),
         "Invert": (lambda num_bins, height, width: None, False),
@@ -324,7 +325,7 @@ class RandAugment(_AutoAugmentBase):
             lambda num_bins, height, width: (8 - (torch.arange(num_bins) / ((num_bins - 1) / 4))).round().int(),
             False,
         ),
-        "Solarize": (lambda num_bins, height, width: torch.linspace(255.0, 0.0, num_bins), False),
+        "Solarize": (lambda num_bins, height, width: torch.linspace(1.0, 0.0, num_bins), False),
         "AutoContrast": (lambda num_bins, height, width: None, False),
         "Equalize": (lambda num_bins, height, width: None, False),
     }
@@ -378,7 +379,7 @@ class TrivialAugmentWide(_AutoAugmentBase):
             lambda num_bins, height, width: (8 - (torch.arange(num_bins) / ((num_bins - 1) / 6))).round().int(),
             False,
         ),
-        "Solarize": (lambda num_bins, height, width: torch.linspace(255.0, 0.0, num_bins), False),
+        "Solarize": (lambda num_bins, height, width: torch.linspace(1.0, 0.0, num_bins), False),
         "AutoContrast": (lambda num_bins, height, width: None, False),
         "Equalize": (lambda num_bins, height, width: None, False),
     }
@@ -423,7 +424,7 @@ class AugMix(_AutoAugmentBase):
             lambda num_bins, height, width: (4 - (torch.arange(num_bins) / ((num_bins - 1) / 4))).round().int(),
             False,
         ),
-        "Solarize": (lambda num_bins, height, width: torch.linspace(255.0, 0.0, num_bins), False),
+        "Solarize": (lambda num_bins, height, width: torch.linspace(1.0, 0.0, num_bins), False),
         "AutoContrast": (lambda num_bins, height, width: None, False),
         "Equalize": (lambda num_bins, height, width: None, False),
     }
@@ -505,13 +506,7 @@ def forward(self, *inputs: Any) -> Any:
                 aug = self._apply_image_or_video_transform(
                     aug, transform_id, magnitude, interpolation=self.interpolation, fill=self.fill
                 )
-            mix.add_(
-                # The multiplication below could become in-place provided `aug is not batch and aug.is_floating_point()`
-                # Currently we can't do this because `aug` has to be `unint8` to support ops like `equalize`.
-                # TODO: change this once all ops in `F` support floats. https://github.com/pytorch/vision/issues/6840
-                combined_weights[:, i].reshape(batch_dims)
-                * aug
-            )
+            mix.add_(combined_weights[:, i].reshape(batch_dims) * aug)
         mix = mix.reshape(orig_dims).to(dtype=image_or_video.dtype)
 
         if isinstance(orig_image_or_video, (features.Image, features.Video)):


Test patch: