Repository: kornia/kornia
Pull request: Implement fast version of crop function in augmentations
### Description

Fixes #966 

### Status
**Ready/Work in progress/Hold**

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not applicable items -->
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] Breaking change (fix or new feature that would cause existing functionality to change)
- [ ] New tests added to cover the changes
- [ ] Docstrings/Documentation updated


## PR Checklist
### PR Implementer
This is a small checklist for the implementation details of this PR.

If there are any questions regarding code style or other conventions check out our 
[summary](https://github.com/kornia/kornia/blob/master/CONTRIBUTING.rst).

- [ ] Did you discuss the functionality or any breaking changes before ?
- [ ] **Pass all tests**: did you test in local ? `make test`
- [ ] Unittests: did you add tests for your new functionality ?
- [ ] Documentations: did you build documentation ? `make build-docs`
- [ ] Implementation: is your code well commented and follow conventions ? `make lint`
- [ ] Docstrings & Typing: has your code documentation and typing ? `make mypy`
- [ ] Update notebooks & documentation if necessary

### KorniaTeam
<details>
  <summary>KorniaTeam workflow</summary>
  
  - [ ] Assign correct label
  - [ ] Assign PR to a reviewer
  - [ ] Does this PR close an Issue? (add `closes #IssueNumber` at the bottom if 
        not already in description)

</details>

### Reviewer
<details>
  <summary>Reviewer workflow</summary>

  - [ ] Do all tests pass? (Unittests, Typing, Linting, Documentation, Environment)
  - [ ] Does the implementation follow `kornia` design conventions?
  - [ ] Is the documentation complete enough ?
  - [ ] Are the tests covering simple and corner cases ?
 
</details>


Problem statement:
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,

Hints from issue 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

Commit message: implement indexing random resize crop

Patch to review:
diff --git a/kornia/augmentation/augmentation.py b/kornia/augmentation/augmentation.py
index 88e9f98c97..e5124c924e 100644
--- a/kornia/augmentation/augmentation.py
+++ b/kornia/augmentation/augmentation.py
@@ -846,11 +846,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,
+        mode: str = 'slice',
     ) -> None:
         # Since PyTorch does not support ragged tensor. So cropping function happens all the time.
         super(RandomResizedCrop, self).__init__(
@@ -863,7 +869,8 @@ 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),
+            mode=mode
         )
 
     def __repr__(self) -> str:
@@ -884,7 +891,26 @@ def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Te
         return F.compute_crop_transformation(input, params, self.flags)
 
     def apply_transform(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
-        return F.apply_crop(input, params, self.flags)
+        if self.flags['mode'] == 'resample':
+            return F.apply_crop(input, params, self.flags)
+        elif self.flags['mode'] == 'slice':
+            import kornia as K
+            coords = params['src'].long()
+            B, C, _, _ = input.shape
+            out = torch.empty(B, C, *self.size, device=input.device, dtype=input.dtype)
+            for i in range(B):
+                x1 = int(coords[i, 0, 0])
+                x2 = int(coords[i, 1, 0]) + 1
+                y1 = int(coords[i, 0, 1])
+                y2 = int(coords[i, 3, 1]) + 1
+                out[i] = K.geometry.transform.resize(
+                    input[i:i+1, :, y1:y2, x1:x2],
+                    self.size,
+                    interpolation='bilinear',
+                    align_corners=True)
+            return out
+        else:
+            raise NotImplementedError(f"Not supported type: {self.flags['mode']}.")
 
 
 class Normalize(AugmentationBase2D):


Code review comments:
diff_hunk: @@ -884,7 +891,26 @@ def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Te
         return F.compute_crop_transformation(input, params, self.flags)
 
     def apply_transform(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
-        return F.apply_crop(input, params, self.flags)
+        if self.flags['mode'] == 'resample':
+            return F.apply_crop(input, params, self.flags)
+        elif self.flags['mode'] == 'slice':
+            import kornia as K
+            coords = params['src'].long()
+            B, C, _, _ = input.shape
+            out = torch.empty(B, C, *self.size, device=input.device, dtype=input.dtype)
+            for i in range(B):
+                x1 = int(coords[i, 0, 0])
+                x2 = int(coords[i, 1, 0]) + 1
+                y1 = int(coords[i, 0, 1])
+                y2 = int(coords[i, 3, 1]) + 1
+                out[i] = K.geometry.transform.resize(
+                    input[i:i+1, :, y1:y2, x1:x2],
+                    self.size,
+                    interpolation='bilinear',
+                    align_corners=True)
+            return out
+        else:
+            raise NotImplementedError(f"Not supported type: {self.flags['mode']}.")
path: kornia/augmentation/augmentation.py
text: @author:
@user1 check if something like this would work.
And for augmentations I would give that one by the default so that gives better performance 

@user1:
Add tests. We are fine with it if it passed all the tests. 

@user1:
Also, can you provide a benchmark between the two implementations?