Repository: kornia/kornia

Issue:
Make contrib.extract_tensor_patches better
## 🚀 Feature

Make `contrib.extract_tensor_patches` "better" (see the conclusion at the bottom for what this means in detail) by replacing the [convolution](https://github.com/kornia/kornia/blob/86185fbcec030565c742c325b64dcb94700b0834/kornia/contrib/extract_patches.py#L109) with [`torch.Tensor.unfold`](https://pytorch.org/docs/stable/tensors.html?highlight=unfold#torch.Tensor.unfold).

## Pitch

Consider the following setup

```python
from kornia.contrib import extract_tensor_patches as extract_tensor_patches_conv


image = torch.randn(2, 3, 1024, 1024)

window_size = 3
stride = 2

patches_conv = extract_tensor_patches_conv(image, window_size, stride=stride)
```

Here is a drop in replacement for `extract_tensor_patches` using `unfold`

```python
from typing import Sequence, Union

import torch
from torch.nn.modules.utils import _pair

def extract_tensor_patchesnd(
    input: torch.Tensor, window_sizes: Sequence[int], strides: Sequence[int]
) -> torch.Tensor:
    batch_size, num_channels = input.size()[:2]
    dims = range(2, input.dim())
    for dim, patch_size, stride in zip(dims, window_sizes, strides):
        input = input.unfold(dim, patch_size, stride)
    input = input.permute(0, *dims, 1, *[dim + len(dims) for dim in dims]).contiguous()
    return input.view(batch_size, -1, num_channels, *window_sizes)


def extract_tensor_patches_unfold(
    input: torch.Tensor,
    window_size: Union[int, Sequence[int]],
    stride: Union[int, Sequence[int]] = 1,
) -> torch.Tensor:
    return extract_tensor_patchesnd(input, _pair(window_size), _pair(stride))

patches_unfold = extract_tensor_patches_unfold(image, window_size, stride=stride)
assert torch.norm(patches_unfold - patches_conv) < 1e-6
```

On the CPU this achieves a speedup of approximately factor `1.8`:

```python
def timeit(fn, image, number=100):
    total = _timeit.timeit(
        lambda: fn(image, window_size, stride=stride), number=number,
    )
    avg = total / number
    if avg < 1e-3:
        factor = 1e6
        unit = "µs"
    else:
        factor = 1e3
        unit = "ms"
    print(f"Average runtime for {fn.__name__}: {avg * factor:.1f} {unit}")


timeit(extract_tensor_patches_conv, image)
timeit(extract_tensor_patches_unfold, image)
```

```
Average runtime for extract_tensor_patches: 70.9 ms
Average runtime for extract_tensor_patches_unfold: 39.2 ms
```

With CUDA this achieves a speedup of approximately factor `66`:

```python
timeit(extract_tensor_patches_conv, image.cuda())
timeit(extract_tensor_patches_unfold, image.cuda())
```

```
Average runtime for extract_tensor_patches: 3.3 ms
Average runtime for extract_tensor_patches_unfold: 50.1 µs
```

## Conclusion

The proposed approach

- is a drop-in replacement and thus should not break BC in any way
- is faster on CPU and much faster on CUDA
- has less custom code to maintain and **IMO** more clear overall
- can be easily extended for other-dimensional inputs

I you want this let me know and I'll send a PR.


Discussion:
@pmeier looks good to me. Go ahead with a PR

Resolving patch:
diff --git a/kornia/contrib/extract_patches.py b/kornia/contrib/extract_patches.py
--- a/kornia/contrib/extract_patches.py
+++ b/kornia/contrib/extract_patches.py
@@ -10,10 +10,6 @@
 class ExtractTensorPatches(nn.Module):
     r"""Module that extract patches from tensors and stack them.
 
-    Applies a 2D convolution over an input tensor to extract patches and stack
-    them in the depth axis of the output tensor. The function applies a
-    Depthwise Convolution by applying the same kernel for all the input planes.
-
     In the simplest case, the output value of the operator with input size
     :math:`(B, C, H, W)` is :math:`(B, N, C, H_{out}, W_{out})`.
 
@@ -42,10 +38,10 @@ class ExtractTensorPatches(nn.Module):
           the height dimension, and the second `int` for the width dimension.
 
     Arguments:
-        window_size (Union[int, Tuple[int, int]]): the size of the convolving
-          kernel and the output patch size.
+        window_size (Union[int, Tuple[int, int]]): the size of the sliding
+          window and the output patch size.
         stride (Optional[Union[int, Tuple[int, int]]]): stride of the
-          convolution. Default is 1.
+          sliding window. Default is 1.
         padding (Optional[Union[int, Tuple[int, int]]]): Zero-padding added to
           both side of the input. Default is 0.
 
@@ -78,59 +74,44 @@ def __init__(
         self.stride: Tuple[int, int] = _pair(stride)
         self.padding: Tuple[int, int] = _pair(padding)
 
-        # create base kernel
-        self.kernel: torch.Tensor = self.create_kernel(self.window_size)
-
-    @staticmethod
-    def create_kernel(
-            window_size: Tuple[int, int]) -> torch.Tensor:
-        r"""Creates a binary kernel to extract the patches. If the window size
-        is HxW will create a (H*W)xHxW kernel.
-        """
-        window_range: int = window_size[0] * window_size[1]
-        kernel: torch.Tensor = torch.zeros(window_range, window_range)
-        for i in range(window_range):
-            kernel[i, i] += 1.0
-        return kernel.view(window_range, 1, window_size[0], window_size[1])
-
     def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore
-        if not torch.is_tensor(input):
-            raise TypeError("Input input type is not a torch.Tensor. Got {}"
-                            .format(type(input)))
-        if not len(input.shape) == 4:
-            raise ValueError("Invalid input shape, we expect BxCxHxW. Got: {}"
-                             .format(input.shape))
-        # unpack shapes
-        batch_size, channels, height, width = input.shape
-
-        # run convolution 2d to extract patches
-        kernel: torch.Tensor = self.kernel.repeat(channels, 1, 1, 1)
-        kernel = kernel.to(input.device).to(input.dtype)
-        output_tmp: torch.Tensor = F.conv2d(
-            input,
-            kernel,
-            stride=self.stride,
-            padding=self.padding,
-            groups=channels)
-
-        # reshape the output tensor
-        output: torch.Tensor = output_tmp.view(
-            batch_size, channels, self.window_size[0], self.window_size[1], -1)
-        return output.permute(0, 4, 1, 2, 3)  # BxNxCxhxw
+        return extract_tensor_patches(input, self.window_size, stride=self.stride, padding=self.padding)
 
 
 ######################
 # functional interface
 ######################
 
+def _extract_tensor_patchesnd(
+    input: torch.Tensor, window_sizes: Tuple[int, ...], strides: Tuple[int, ...]
+) -> torch.Tensor:
+    batch_size, num_channels = input.size()[:2]
+    dims = range(2, input.dim())
+    for dim, patch_size, stride in zip(dims, window_sizes, strides):
+        input = input.unfold(dim, patch_size, stride)
+    input = input.permute(0, *dims, 1, *[dim + len(dims) for dim in dims]).contiguous()
+    return input.view(batch_size, -1, num_channels, *window_sizes)
+
 
 def extract_tensor_patches(
-        input: torch.Tensor,
-        window_size: Union[int, Tuple[int, int]],
-        stride: Optional[Union[int, Tuple[int, int]]] = 1,
-        padding: Optional[Union[int, Tuple[int, int]]] = 0) -> torch.Tensor:
+    input: torch.Tensor,
+    window_size: Union[int, Tuple[int, int]],
+    stride: Union[int, Tuple[int, int]] = 1,
+    padding: Union[int, Tuple[int, int]] = 0
+) -> torch.Tensor:
     r"""Function that extract patches from tensors and stack them.
 
     See :class:`~kornia.contrib.ExtractTensorPatches` for details.
     """
-    return ExtractTensorPatches(window_size, stride, padding)(input)
+    if not torch.is_tensor(input):
+        raise TypeError("Input input type is not a torch.Tensor. Got {}"
+                        .format(type(input)))
+    if not len(input.shape) == 4:
+        raise ValueError("Invalid input shape, we expect BxCxHxW. Got: {}"
+                         .format(input.shape))
+
+    if padding:
+        pad_vert, pad_horz = _pair(padding)
+        input = F.pad(input, [pad_horz, pad_horz, pad_vert, pad_vert])
+
+    return _extract_tensor_patchesnd(input, _pair(window_size), _pair(stride))


Test patch: