Rearrange torch 2D tensors ("Tiles") to be in a particular order

user_10969548 0
Viewed 1264

I have something that looks like this...

import torch
X = torch.cat([torch.ones((2,4)).unsqueeze(0), torch.ones((2,4)).unsqueeze(0)*2, torch.ones((2,4)).unsqueeze(0)*3, torch.ones((2,4)).unsqueeze(0)*4]).unsqueeze(0)

which is

tensor([[[[1., 1., 1., 1.],
          [1., 1., 1., 1.]],

         [[2., 2., 2., 2.],
          [2., 2., 2., 2.]],

         [[3., 3., 3., 3.],
          [3., 3., 3., 3.]],

         [[4., 4., 4., 4.],
          [4., 4., 4., 4.]]]])

And I want to rearrange them to look like this

tensor([[[1., 1., 1., 1., 2., 2., 2., 2.],
         [1., 1., 1., 1., 2., 2., 2., 2.],
         [3., 3., 3., 3., 4., 4., 4., 4.],
         [3., 3., 3., 3., 4., 4., 4., 4.]]])

However the solution I was hoping would work obviously doesn't and while I understand why it doesn't work I'll share it.

b, t, x, y = X.shape
num_x_splits = num_y_splits = 2
X = X.reshape(-1,x*num_x_splits,y*num_y_splits)

tensor([[[1., 1., 1., 1., 1., 1., 1., 1.],
         [2., 2., 2., 2., 2., 2., 2., 2.],
         [3., 3., 3., 3., 3., 3., 3., 3.],
         [4., 4., 4., 4., 4., 4., 4., 4.]]])

Is there a computationally efficient way I can accomplish what I want?

I did some more research and looks like the ideal solution might actually be a combinatoin of the accepted answer and to use the einops package which provides a high-level, easy-to-read, efficient solution

pytorch numpy python
5 Answers

You asked for einops one-liner? Here you go:

>>> from einops import rearrange

>>> x.shape  # x is taken from question.
torch.Size([1, 4, 2, 4])

>>> rearrange(x, '1 (i1 i2) j k -> (i1 j) (i2 k)', i1=2)
tensor([[1., 1., 1., 1., 2., 2., 2., 2.],
        [1., 1., 1., 1., 2., 2., 2., 2.],
        [3., 3., 3., 3., 4., 4., 4., 4.],
        [3., 3., 3., 3., 4., 4., 4., 4.]])
user_4988920

I will start with the simple case where we know the number of tensors and which tensors to be positioned etc.:

ones = 1 * torch.ones(2,4)
twos = 2 * torch.ones(2,4)
thrs = 3 * torch.ones(2,4)
fors = 4 * torch.ones(2,4)
fivs = 5 * torch.ones(2,4)
sixs = 6 * torch.ones(2,4)
row1 = torch.cat([ones, twos], axis = 1)
row2 = torch.cat([thrs, fors], axis = 1)
row3 = torch.cat([fivs, sixs], axis = 1)
comb = torch.cat([row1, row2, row3], axis = 0)
print(comb)

tensor([[1., 1., 1., 1., 2., 2., 2., 2.],
        [1., 1., 1., 1., 2., 2., 2., 2.],
        [3., 3., 3., 3., 4., 4., 4., 4.],
        [3., 3., 3., 3., 4., 4., 4., 4.],
        [5., 5., 5., 5., 6., 6., 6., 6.],
        [5., 5., 5., 5., 6., 6., 6., 6.]])

So you just need to concatenate rows first then combine all rows on the axis=0 which will get the positioning as you want if you know which tensors you want to combine etc.

Another way to solve this is to permute the dimensions so that you get the placement you want which should work for more complex cases as well:

comb2 = torch.cat([ones, twos, thrs, fors, fivs, sixs]).view(3, 2, 2, 4)
print(comb2)
print(comb2.permute(0, 2, 1, 3). reshape(6, 8))

tensor([[[[1., 1., 1., 1.],
          [1., 1., 1., 1.]],

         [[2., 2., 2., 2.],
          [2., 2., 2., 2.]]],


        [[[3., 3., 3., 3.],
          [3., 3., 3., 3.]],

         [[4., 4., 4., 4.],
          [4., 4., 4., 4.]]],


        [[[5., 5., 5., 5.],
          [5., 5., 5., 5.]],

         [[6., 6., 6., 6.],
          [6., 6., 6., 6.]]]])
tensor([[1., 1., 1., 1., 2., 2., 2., 2.],
        [1., 1., 1., 1., 2., 2., 2., 2.],
        [3., 3., 3., 3., 4., 4., 4., 4.],
        [3., 3., 3., 3., 4., 4., 4., 4.],
        [5., 5., 5., 5., 6., 6., 6., 6.],
        [5., 5., 5., 5., 6., 6., 6., 6.]])

In the first view of the tensor, we view it such that we know we have 3 sets of tensors with 2x4 dimensions and each of these 2x4 tensors stored in additional dimension (2x2x4) to make it easier to manipulate instead of storing 2x8 or 4x4. Then we rearrange the dimensions to get the placement you want. I try to go backwards from the desired dimension to find the permutation I want usually. First dimension should then be 6 = 3x2 where we get 3 sets and 2 rows of tensor so we keep the first axis in place, move rows dimension next to set dimension: permute(0, 2, 1, 3). Permutation of 4 dimensional tensors can be a bit tricky to imagine, just play with it and you will get the right combination :)

user_16071040

I accepted the correct answer but just wanted to share a more generic version of the code in case it helps anyone.

x, y = tile_x_dim, tile_y_dim   # for readability
X = X.reshape(-1, num_x_splits, num_y_splits, x, y)
X = X.permute(0,1,3,2,4).reshape(-1,x*num_x_splits,y*num_y_splits)

Where tile_x_dim and tile_y_dim are the x and y dimension respectively of the tile that was formed from the full image. num_x_splits and num_y_splits is the number of times the original picture was split along the x-axis. Note that num_x_splits multiplied by num_y_splits would give the number of tiles the image was cut into.

user_109695480

You can use NumPy's block() function to create block-matrices (and then convert to a tensor):

import torch
import numpy as np

A = np.ones((2,4))
B = np.ones((2,4))*2
C = np.ones((2,4))*3
D = np.ones((2,4))*4

X = np.block([[A,B],
              [C,D]])

torch.tensor(X)

tensor([[1., 1., 1., 1., 2., 2., 2., 2.],
        [1., 1., 1., 1., 2., 2., 2., 2.],
        [3., 3., 3., 3., 4., 4., 4., 4.],
        [3., 3., 3., 3., 4., 4., 4., 4.]], dtype=torch.float64)
user_90676150 
import torch
X = torch.cat([torch.ones((2,4)).unsqueeze(0), torch.ones((2,4)).unsqueeze(0)*3], dim = 1)
Y = torch.cat([torch.ones((2,4)).unsqueeze(0)*2, torch.ones((2,4)).unsqueeze(0)*4], dim = 1)
Z = torch.cat((X, Y), dim = 2)
print(Z)

tensor([[[1., 1., 1., 1., 2., 2., 2., 2.],
         [1., 1., 1., 1., 2., 2., 2., 2.],
         [3., 3., 3., 3., 4., 4., 4., 4.],
         [3., 3., 3., 3., 4., 4., 4., 4.]]])

Generally we avoid tensor to numpy conversions when we needed to use the tensor for Back-propagating. I think you can use this and its the exact dimensions you are looking for.

user_94894650
Related