Repository: pytorch/vision

Issue:
Fast rotation for 90, 180, and 270 degree
### 🚀 The feature

90, 180, and 270 degree rotation can be done without interpolation by simply re-arranging the data. This is significantly faster, especially for high resolution image.

### Motivation, pitch

Image rotation becomes a bottleneck for me when using large images (1000x1000). Currently PyTorch already has `torch.rot90()`, which is significantly faster than `torchvision.transforms.v2.functional.rotate()`. In some applications, it is acceptable to use only 90, 180, and 270 degree rotation for augmentation (cross-ref #566).

Perhaps `torchvision.transforms.v2.functional.rotate()` can check for 90, 180, and 270 specifically, and call `torch.rot90()`?

### Alternatives

A workaround is to write a custom transform to use `torch.rot90()`.

### Additional context

Micro-benchmark results

```python
import torch
import torchvision.transforms.v2.functional as TF
import time

N = 100

input = torch.randn(3, 1024, 1024)
output1 = TF.rotate(input, 90)
output2 = torch.rot90(input, k=1, dims=(1, 2))

torch.testing.assert_close(output1, output2)

def benchmark(fn, *args):
    time0 = time.perf_counter()
    for _ in range(N):
        fn(input, *args)
    return N / (time.perf_counter() - time0)

print(benchmark(TF.rotate, 90))
print(benchmark(torch.rot90, 1, (1, 2)))
```

Function | Throughput (it/s)
------------|-----------------------
`TF.rotate()` | 52.95
`torch.rot90()` | 10,748.12


Discussion:
Thank you for the feature request @gau-nernst . This is indeed something we should consider for inclusion. Let us know if you're interested in submitting a PR for this.
@NicolasHug Sure, I can give it a try.
Thanks! Let's try to add that functionality only to `torchvision.transforms.v2.functional.rotate` (i.e. you do not need to worry about `torchvision.transforms.functional.rotate`)

Resolving patch:
diff --git a/torchvision/transforms/v2/functional/_geometry.py b/torchvision/transforms/v2/functional/_geometry.py
--- a/torchvision/transforms/v2/functional/_geometry.py
+++ b/torchvision/transforms/v2/functional/_geometry.py
@@ -997,6 +997,21 @@ def rotate_image(
     center: Optional[List[float]] = None,
     fill: _FillTypeJIT = None,
 ) -> torch.Tensor:
+    angle = angle % 360  # shift angle to [0, 360) range
+
+    # fast path: transpose without affine transform
+    if center is None:
+        if angle == 0:
+            return image.clone()
+        if angle == 180:
+            return torch.rot90(image, k=2, dims=(-2, -1))
+
+        if expand or image.shape[-1] == image.shape[-2]:
+            if angle == 90:
+                return torch.rot90(image, k=1, dims=(-2, -1))
+            if angle == 270:
+                return torch.rot90(image, k=3, dims=(-2, -1))
+
     interpolation = _check_interpolation(interpolation)
 
     input_height, input_width = image.shape[-2:]


Test patch:
diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py
--- a/test/test_transforms_v2.py
+++ b/test/test_transforms_v2.py
@@ -1782,6 +1782,17 @@ def test_transform_unknown_fill_error(self):
         with pytest.raises(TypeError, match="Got inappropriate fill arg"):
             transforms.RandomAffine(degrees=0, fill="fill")
 
+    @pytest.mark.parametrize("size", [(11, 17), (16, 16)])
+    @pytest.mark.parametrize("angle", [0, 90, 180, 270])
+    @pytest.mark.parametrize("expand", [False, True])
+    def test_functional_image_fast_path_correctness(self, size, angle, expand):
+        image = make_image(size, dtype=torch.uint8, device="cpu")
+
+        actual = F.rotate(image, angle=angle, expand=expand)
+        expected = F.to_image(F.rotate(F.to_pil_image(image), angle=angle, expand=expand))
+
+        torch.testing.assert_close(actual, expected)
+
 
 class TestContainerTransforms:
     class BuiltinTransform(transforms.Transform):