Repository: kornia/kornia

Issue:
Slower than new PyTorch GPU transformations
## ❓ Questions and Help
Kornia's augmentations seem to be slower than PyTorch's one. The dataset I am using is CelebA.

I have:

OS: Ubuntu 20.04
GPU: 1080ti
PyTorch: 1.8.0
torchvision: 0.9.0

```python

from torchvision.io import read_image
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import torchvision.transforms as TT
import kornia.augmentation as KT

def timeit(func):
    def _inner(*args, **kwargs):
        start = time()
        func(*args, **kwargs)
        end = time()
        print(f"Elapsed = {end - start :.2f}s")

    return _inner

@timeit
def test_augmentation(tr, device="cuda"):
    ds = ImageFolder(root="~/Documents/Datasets/celeba", loader=read_image)
    dl = DataLoader(ds, batch_size=128, num_workers=16, pin_memory=True)

    for batch in tqdm(dl):
        x = batch[0].to(device)
        tr(x)
        continue


def test_augmentations():
    resize = 256
    crop = 224
    device = "cuda"

    augmentations = [
        nn.Sequential(
            TT.Resize((resize, resize)),
            TT.ConvertImageDtype(dtype=torch.float),
            TT.RandomVerticalFlip(),
            TT.RandomHorizontalFlip(),
            TT.RandomCrop(crop),
        ).to(device),
        nn.Sequential(
            TT.Resize((resize, resize)),
            TT.ConvertImageDtype(dtype=torch.float),
            KT.RandomVerticalFlip(),
            KT.RandomHorizontalFlip(),
            KT.RandomCrop((crop, crop)),
        ).to(device),
    ]

    list(map(test_augmentation, augmentations))


test_augmentations()
```
PyTorch's augmentations takes 15s while Kornia takes 29s. 

Best regards,


Discussion:
Hi @FrancescoSaverioZuppichini,

Thank you for reporting. Torchvision implemented the RandomCrop and related in a non-differentiable way, while Kornia is not. 

We are aware of the timing issue that brought by the differentiable cropping implementation, and we are planning on incorporating the non-differentiable implementation for those common usages. 

Best,
Jian
Besides what @shijianjian pointed about  differentiability, kornia is meant to work better in batch on device which means adding kornia.augmentation as a layer before your model instead of inside the dataloader.

In addition, our plan is to not get fully aligned with torchvision augmentations since our use cases and future APIs might be different.
Hi @shijianjian, thank you for the quick reply. 

> Besides what @shijianjian pointed about differentiability, kornia is meant to work better in batch on device which means adding kornia.augmentation as a layer before your model instead of inside the dataloader.

That's the same thing, It won't change performances. I am running them on the batched data on the GPU.

I am felling pretty dumb today, but correct me if I am wrong. The following snippet:

```python
import torchvision.transforms as TT

resize = 256
crop = 224
device = "cuda"
x = torch.randn((1, 3, 224, 224), requires_grad=True)
y = torch.ones_like(x)
# torchvision augmentation
tr = nn.Sequential(TT.RandomVerticalFlip())
model = nn.Conv2d(3, 3, kernel_size=1)
criterion = nn.MSELoss()
optimizer = opt.Adam(model.parameters())
for _ in range(100):
    out = tr(model(x))
    loss = criterion(y, out)
    optimizer.zero_grad()
    loss.backward()
    print(loss)
    optimizer.step()
```

Shouldn't work since `torchvision.augmentation` is not differentiable, but it works. I think I am missing something. Nevertheless, thank you!
Non-differentiable does not necessarily mean to break the computational graph. As a differentiable operation, apart from retaining the computational graph, it will also provide meaningful gradients. e.g. the crop function in kornia is based on perspective transform which provides the gradients w.r.t. the grid.

In your example, the gradients will be retained since the graph is still intact. Under the hood, TorchVision and Kornia implemented the flip in the same way. For cropping and related, kornia will provide traceable gradients that differs from torchvision.
Thank you for the clarifications

Resolving patch:
diff --git a/kornia/augmentation/augmentation.py b/kornia/augmentation/augmentation.py
--- a/kornia/augmentation/augmentation.py
+++ b/kornia/augmentation/augmentation.py
@@ -38,6 +38,7 @@
 )
 from kornia.utils import _extract_device_dtype
 from kornia.enhance.normalize import normalize, denormalize
+from kornia.geometry import resize
 
 from . import random_generator as rg
 from .utils import (
@@ -619,6 +620,10 @@ class CenterCrop(AugmentationBase2D):
             applied to each. Default: False.
         keepdim (bool): whether to keep the output shape the same as input (True) or broadcast it
                         to the batch form (False). Default: False.
+        cropping_mode (str): The used algorithm to crop. ``slice`` will use advanced slicing to extract the tensor based
+                             on the sampled indices. ``resample`` will use `warp_affine` using the affine transformation
+                             to extract and resize at once. Use `slice` for efficiency, or `resample` for proper
+                             differentiability. Default: `slice`.
 
     Shape:
         - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
@@ -643,9 +648,16 @@ class CenterCrop(AugmentationBase2D):
                   [-1.2633,  0.3500]]]])
     """
 
-    def __init__(self, size: Union[int, Tuple[int, int]], align_corners: bool = True,
-                 resample: Union[str, int, Resample] = Resample.BILINEAR.name,
-                 return_transform: bool = False, p: float = 1., keepdim: bool = False) -> None:
+    def __init__(
+            self,
+            size: Union[int, Tuple[int, int]],
+            align_corners: bool = True,
+            resample: Union[str, int, Resample] = Resample.BILINEAR.name,
+            return_transform: bool = False,
+            p: float = 1.,
+            keepdim: bool = False,
+            cropping_mode: str = 'slice',
+    ) -> None:
         # same_on_batch is always True for CenterCrop
         # Since PyTorch does not support ragged tensor. So cropping function happens batch-wisely.
         super(CenterCrop, self).__init__(p=1., return_transform=return_transform, same_on_batch=True, p_batch=p,
@@ -663,6 +675,7 @@ def __init__(self, size: Union[int, Tuple[int, int]], align_corners: bool = True
             interpolation=torch.tensor(self.resample.value),
             align_corners=torch.tensor(align_corners)
         )
+        self.cropping_mode = cropping_mode
 
     def __repr__(self) -> str:
         repr = f"size={self.size}"
@@ -681,9 +694,24 @@ def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Te
     def apply_transform(
         self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
-        transform = cast(torch.Tensor, transform)
-        return crop_by_transform_mat(
-            input, transform[:, :2, :], self.size, self.resample.name.lower(), 'zeros', self.align_corners)
+        if self.cropping_mode == 'resample':  # uses bilinear interpolation to crop
+            transform = cast(torch.Tensor, transform)
+            return crop_by_transform_mat(
+                input, transform[:, :2, :], self.size, self.resample.name.lower(), 'zeros', self.align_corners)
+        elif self.cropping_mode == 'slice':   # uses advanced slicing to crop
+            # TODO: implement as separated function `crop_and_resize_iterative`
+            B, C, _, _ = input.shape
+            H, W = self.size
+            out = torch.empty(B, C, H, W, device=input.device, dtype=input.dtype)
+            for i in range(B):
+                x1 = int(params['src'][i, 0, 0])
+                x2 = int(params['src'][i, 1, 0]) + 1
+                y1 = int(params['src'][i, 0, 1])
+                y2 = int(params['src'][i, 3, 1]) + 1
+                out[i] = input[i:i + 1, :, y1:y2, x1:x2]
+            return out
+        else:
+            raise NotImplementedError(f"Not supported type: {self.cropping_mode}.")
 
 
 class RandomRotation(AugmentationBase2D):
@@ -798,9 +826,13 @@ class RandomCrop(AugmentationBase2D):
                                       input tensor. If ``False`` and the input is a tuple the applied transformation
                                       wont be concatenated
         same_on_batch (bool): apply the same transformation across the batch. Default: False
-        align_corners(bool): interpolation flag. Default: False.
+        align_corners(bool): interpolation flag. Default: True.
         keepdim (bool): whether to keep the output shape the same as input (True) or broadcast it
                         to the batch form (False). Default: False.
+        cropping_mode (str): The used algorithm to crop. ``slice`` will use advanced slicing to extract the tensor based
+                             on the sampled indices. ``resample`` will use `warp_affine` using the affine transformation
+                             to extract and resize at once. Use `slice` for efficiency, or `resample` for proper
+                             differentiability. Default: `slice`.
 
     Shape:
         - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
@@ -812,20 +844,28 @@ class RandomCrop(AugmentationBase2D):
         applied transformation will be merged int to the input transformation tensor and returned.
 
     Examples:
-        >>> rng = torch.manual_seed(0)
-        >>> inputs = torch.randn(1, 1, 3, 3)
+        >>> _ = torch.manual_seed(0)
+        >>> inputs = torch.arange(1*1*3*3.).view(1, 1, 3, 3)
         >>> aug = RandomCrop((2, 2), p=1.)
         >>> aug(inputs)
-        tensor([[[[-0.6562, -1.0009],
-                  [ 0.2223, -0.5507]]]])
+        tensor([[[[3., 4.],
+                  [6., 7.]]]])
     """
 
     def __init__(
-        self, size: Tuple[int, int], padding: Optional[Union[int, Tuple[int, int], Tuple[int, int, int, int]]] = None,
-        pad_if_needed: Optional[bool] = False, fill: int = 0, padding_mode: str = 'constant',
+        self,
+        size: Tuple[int, int],
+        padding: Optional[Union[int, Tuple[int, int], Tuple[int, int, int, int]]] = None,
+        pad_if_needed: Optional[bool] = False,
+        fill: int = 0,
+        padding_mode: str = 'constant',
         resample: Union[str, int, Resample] = Resample.BILINEAR.name,
-        return_transform: bool = False, same_on_batch: bool = False, align_corners: bool = False, p: float = 1.0,
-        keepdim: bool = False
+        return_transform: bool = False,
+        same_on_batch: bool = False,
+        align_corners: bool = True,
+        p: float = 1.0,
+        keepdim: bool = False,
+        cropping_mode: str = 'slice',
     ) -> None:
         # Since PyTorch does not support ragged tensor. So cropping function happens batch-wisely.
         super(RandomCrop, self).__init__(
@@ -841,6 +881,7 @@ def __init__(
             interpolation=torch.tensor(self.resample.value),
             align_corners=torch.tensor(align_corners)
         )
+        self.cropping_mode = cropping_mode
 
     def __repr__(self) -> str:
         repr = (f"crop_size={self.size}, padding={self.padding}, fill={self.fill}, pad_if_needed={self.pad_if_needed}, "
@@ -883,10 +924,23 @@ def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Te
     def apply_transform(
         self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
-        transform = cast(torch.Tensor, transform)
-        return crop_by_transform_mat(
-            input, transform, self.size, mode=self.resample.name.lower(),
-            padding_mode='zeros', align_corners=self.align_corners)
+        if self.cropping_mode == 'resample':  # uses bilinear interpolation to crop
+            transform = cast(torch.Tensor, transform)
+            return crop_by_transform_mat(
+                input, transform, self.size, mode=self.resample.name.lower(),
+                padding_mode='zeros', align_corners=self.align_corners)
+        elif self.cropping_mode == 'slice':   # uses advanced slicing to crop
+            B, C, _, _ = input.shape
+            out = torch.empty(B, C, *self.size, device=input.device, dtype=input.dtype)
+            for i in range(B):
+                x1 = int(params['src'][i, 0, 0])
+                x2 = int(params['src'][i, 1, 0]) + 1
+                y1 = int(params['src'][i, 0, 1])
+                y2 = int(params['src'][i, 3, 1]) + 1
+                out[i] = input[i:i + 1, :, y1:y2, x1:x2]
+            return out
+        else:
+            raise NotImplementedError(f"Not supported type: {self.flags['mode']}.")
 
     def forward(self, input: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                 params: Optional[Dict[str, torch.Tensor]] = None, return_transform: Optional[bool] = None
@@ -926,6 +980,10 @@ class RandomResizedCrop(AugmentationBase2D):
         align_corners(bool): interpolation flag. Default: False.
         keepdim (bool): whether to keep the output shape the same as input (True) or broadcast it
                         to the batch form (False). Default: False.
+        cropping_mode (str): The used algorithm to crop. ``slice`` will use advanced slicing to extract the tensor based
+                             on the sampled indices. ``resample`` will use `warp_affine` using the affine transformation
+                             to extract and resize at once. Use `slice` for efficiency, or `resample` for proper
+                             differentiability. Default: `slice`.
 
     Shape:
         - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
@@ -949,11 +1007,17 @@ class RandomResizedCrop(AugmentationBase2D):
     """
 
     def __init__(
-        self, size: Tuple[int, int], scale: Union[torch.Tensor, Tuple[float, float]] = (0.08, 1.0),
+        self,
+        size: Tuple[int, int],
+        scale: Union[torch.Tensor, Tuple[float, float]] = (0.08, 1.0),
         ratio: Union[torch.Tensor, Tuple[float, float]] = (3. / 4., 4. / 3.),
         resample: Union[str, int, Resample] = Resample.BILINEAR.name,
-        return_transform: bool = False, same_on_batch: bool = False,
-        align_corners: bool = True, p: float = 1., keepdim: bool = False
+        return_transform: bool = False,
+        same_on_batch: bool = False,
+        align_corners: bool = True,
+        p: float = 1.,
+        keepdim: bool = False,
+        cropping_mode: str = 'slice',
     ) -> None:
         # Since PyTorch does not support ragged tensor. So cropping function happens all the time.
         super(RandomResizedCrop, self).__init__(
@@ -966,8 +1030,9 @@ def __init__(
         self.align_corners = align_corners
         self.flags: Dict[str, torch.Tensor] = dict(
             interpolation=torch.tensor(self.resample.value),
-            align_corners=torch.tensor(align_corners)
+            align_corners=torch.tensor(align_corners),
         )
+        self.cropping_mode = cropping_mode
 
     def __repr__(self) -> str:
         repr = f"size={self.size}, scale={self.scale}, ratio={self.ratio}, interpolation={self.resample.name}"
@@ -992,10 +1057,26 @@ def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Te
     def apply_transform(
         self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
-        transform = cast(torch.Tensor, transform)
-        return crop_by_transform_mat(
-            input, transform, self.size,
-            mode=self.resample.name.lower(), padding_mode='zeros', align_corners=self.align_corners)
+        if self.cropping_mode == 'resample':  # uses bilinear interpolation to crop
+            transform = cast(torch.Tensor, transform)
+            return crop_by_transform_mat(
+                input, transform, self.size, mode=self.resample.name.lower(),
+                padding_mode='zeros', align_corners=self.align_corners)
+        elif self.cropping_mode == 'slice':   # uses advanced slicing to crop
+            B, C, _, _ = input.shape
+            out = torch.empty(B, C, *self.size, device=input.device, dtype=input.dtype)
+            for i in range(B):
+                x1 = int(params['src'][i, 0, 0])
+                x2 = int(params['src'][i, 1, 0]) + 1
+                y1 = int(params['src'][i, 0, 1])
+                y2 = int(params['src'][i, 3, 1]) + 1
+                out[i] = resize(input[i:i + 1, :, y1:y2, x1:x2],
+                                self.size,
+                                interpolation=(self.resample.name).lower(),
+                                align_corners=self.align_corners)
+            return out
+        else:
+            raise NotImplementedError(f"Not supported type: {self.cropping_mode}.")
 
 
 class Normalize(AugmentationBase2D):


Test patch:
diff --git a/test/augmentation/test_augmentation.py b/test/augmentation/test_augmentation.py
--- a/test/augmentation/test_augmentation.py
+++ b/test/augmentation/test_augmentation.py
@@ -1780,6 +1780,17 @@ def test_no_transform_tuple(self, device, dtype):
         out = kornia.augmentation.CenterCrop((3, 4))(inp)
         assert out.shape == (1, 2, 3, 4)
 
+    def test_crop_modes(self, device, dtype):
+        torch.manual_seed(0)
+        img = torch.rand(1, 3, 5, 5, device=device, dtype=dtype)
+
+        op1 = CenterCrop(size=(2, 2), cropping_mode='resample')
+        out = op1(img)
+
+        op2 = CenterCrop(size=(2, 2), cropping_mode='slice')
+
+        assert_allclose(out, op2(img, op1._params))
+
     def test_gradcheck(self, device, dtype):
         input = torch.rand(1, 2, 3, 4, device=device, dtype=dtype)
         input = utils.tensor_to_gradcheck_var(input)  # to var
@@ -2129,6 +2140,21 @@ def test_pad_if_needed(self, device, dtype):
 
         assert_allclose(out, expected, atol=1e-4, rtol=1e-4)
 
+    def test_crop_modes(self, device, dtype):
+        torch.manual_seed(0)
+        img = torch.tensor([[
+            [0., 1., 2.],
+            [3., 4., 5.],
+            [6., 7., 8.]
+        ]], device=device, dtype=dtype)
+
+        op1 = RandomCrop(size=(2, 2), cropping_mode='resample')
+        out = op1(img)
+
+        op2 = RandomCrop(size=(2, 2), cropping_mode='slice')
+
+        assert_allclose(out, op2(img, op1._params))
+
     def test_gradcheck(self, device, dtype):
         torch.manual_seed(0)  # for random reproductibility
         inp = torch.rand((3, 3, 3), device=device, dtype=dtype)  # 3 x 3
@@ -2263,6 +2289,21 @@ def test_crop_scale_ratio_batch(self, device, dtype):
         out = rrc(inp)
         assert_allclose(out, expected, rtol=1e-4, atol=1e-4)
 
+    def test_crop_modes(self, device, dtype):
+        torch.manual_seed(0)
+        img = torch.tensor([[
+            [0., 1., 2.],
+            [3., 4., 5.],
+            [6., 7., 8.]
+        ]], device=device, dtype=dtype)
+
+        op1 = RandomResizedCrop(size=(4, 4), cropping_mode='resample')
+        out = op1(img)
+
+        op2 = RandomResizedCrop(size=(4, 4), cropping_mode='slice')
+
+        assert_allclose(out, op2(img, op1._params))
+
     def test_gradcheck(self, device, dtype):
         torch.manual_seed(0)  # for random reproductibility
         inp = torch.rand((1, 3, 3), device=device, dtype=dtype)  # 3 x 3
@@ -2274,7 +2315,7 @@ def test_gradcheck(self, device, dtype):
 class TestRandomEqualize:
     # TODO: improve and implement more meaningful smoke tests e.g check for a consistent
     # return values such a torch.Tensor variable.
-    @pytest.mark.xfail(reason="might fail under windows OS due to printing preicision.")
+    @pytest.mark.xfail(reason="might fail under windows OS due to printing precision.")
     def test_smoke(self, device, dtype):
         f = RandomEqualize(p=0.5)
         repr = "RandomEqualize(p=0.5, p_batch=1.0, same_on_batch=False, return_transform=False)"