Strided reshaping of a torch tensor

marukosan 0
Viewed 23

I'm looking for a simple way to reshape a torch tensor x of size (2,8) to a tensor y of size (4,4) so that patches of size (1,4) from x become patches of size (2,2) in y. Here's is the example of what x and y should look like:

x:
[[ 0, 1, 2, 3, 4, 5, 6, 7],
[ 8, 9, 10, 11, 12, 13, 14, 15]]

y:
[[ 0., 1., 4., 5.],
[ 2., 3., 6., 7.],
[ 8., 9., 12., 13.],
[10., 11., 14., 15.]])

I've tried looking into existing torch functions but haven't found any suitable one yet. Is there a simple way to perform this kind of reshaping of a tensor?

torch tensor reshape
1 Answers

Using the amazing einops library (github repo):

from einops import rearrange
y = rearrange(x, 'b (h1 h2 w) -> (b h2) (h1 w)', h1=2, h2=2)

That's it!

It only uses native torch operations as long as x is a torch tensor.

Valentin Goldité0
Related