A vectorized sliding window: clever tensor reshaping/expansion

vision
1

Hi all,

Let’s say I have a tensor, with dims NxHxWxC, for this example 1x4x4x1. I’m trying to do some reshaping of the tensor with a sliding window. This window has a kernel size of k=3, which slides over the W and H dims. I want something like this

t = torch.Tensor([[
[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16],
]])

=>

[[
[1, 2, 3, 5, 6, 7, 9, 10, 11], # i.e. the top-left corner with k=3
[2, 3, 4, 6, 7, 8, 10, 11, 12], # one step to the left
…,
[6,7,8,10,11,12,14,15,16] # bottom-right corner with k=3
]]

The output dim should be N x ? x k **2 x C. If k=3, for this example, it should be 1x6x9x1.

Is this somehow possible, in an efficient way?

Thanks!

2

You can compute it by slicing, flattening and stacking in a simple way. The only drawback is it would be parallelized…

3

This way:

import torch
from torch.nn import functional as f
x = torch.arange(1, 17).view(4, 4).float()
x = x[None, None, :, :]
f.unfold(x, kernel_size=2)

It prints:

tensor([[[ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],
         [ 2.,  3.,  4.,  6.,  7.,  8., 10., 11., 12.],
         [ 5.,  6.,  7.,  9., 10., 11., 13., 14., 15.],
         [ 6.,  7.,  8., 10., 11., 12., 14., 15., 16.]]])
2 Likes
4

Thanks! This is was I needed!