Fast rotation for 90, 180, and 270 degree

[gau-nernst]

🚀 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

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

[NicolasHug]

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.

[gau-nernst]

@NicolasHug Sure, I can give it a try.

[NicolasHug]

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)