"view_as_windows" from skimage but in Pytorch

I needed the same functionality from Pytorch and ended up implementing it myself:

def view_as_windows_torch(image, shape, stride=None):
    """View tensor as overlapping rectangular windows, with a given stride.

    Parameters
    ----------
    image : `~torch.Tensor`
        4D image tensor, with the last two dimensions
        being the image dimensions
    shape : tuple of int
        Shape of the window.
    stride : tuple of int
        Stride of the windows. By default it is half of the window size.

    Returns
    -------
    windows : `~torch.Tensor`
        Tensor of overlapping windows

    """
    if stride is None:
        stride = shape[0] // 2, shape[1] // 2

    windows = image.unfold(2, shape[0], stride[0])
    return windows.unfold(3, shape[1], stride[1])

Essentially it is just two lines of Pytorch code relying on torch.Tensor.unfold. You can easily convince yourself, that it does the same as skimage.util.view_as_windows:

import torch

x = torch.arange(16).reshape((1, 1, 4, 4))

patches = view_as_windows_torch(image=x, shape=(2, 2))
print(patches)

Gives:

tensor([[[[[[ 0,  1],
            [ 4,  5]],

           [[ 1,  2],
            [ 5,  6]],

           [[ 2,  3],
            [ 6,  7]]],


          [[[ 4,  5],
            [ 8,  9]],

           [[ 5,  6],
            [ 9, 10]],

           [[ 6,  7],
            [10, 11]]],


          [[[ 8,  9],
            [12, 13]],

           [[ 9, 10],
            [13, 14]],

           [[10, 11],
            [14, 15]]]]]])

I hope this helps!