Windowed tensor duplication

1

Hi all,

I am trying to reshape and duplicate a tensor based on a window that is applied to one of the
Here is a code example:

x = torch.rand([1, 10, 60, 256]) # [batch, subsamples, timepoints, channels]
window_size = 3

st = torch.arange(0, x.shape[1] - window_size + 1, 1)
nd = st + window_size

t_list = list()
for s, n in zip(st, nd):
    t_list.append(x[:, s:n, :, :].clone())

new_x = torch.stack(t_list).squeeze(1)
new_x.shape() # torch.Size([8, 3, 60, 256])

I was wondering if there is a tensor function that does this, or if there is a more elegant manner to do so, perhaps without using a for loop.

Thanks!

2

You could use tensor.unfold for it:

out = x.unfold(1, window_size, 1)
out = out.squeeze(0).permute(0, 3, 1, 2)
print(out.shape) # torch.Size([8, 3, 60, 256])
print((out == new_x).all()) # tensor(True)
3

Thanks @ptrblck! This is exactly what I was looking for.