Reshape rows into groups of columns

user_3666302 0
Viewed 607

I have a number of row vectors which I would like to batch as column vectors and use as input for Conv1d. As an example I'd like to reshape the tensor x into y i.e. making two groups of two column vectors.

# size = [4, 3]
x = torch.tensor([
    [0,  1,  2],
    [3,  4,  5],
    [6,  7,  8],
    [9, 10, 11]
])

# size = [2, 3, 2]
y = torch.tensor([
    [[0,  3],
     [1,  4],
     [2,  5]],
    [[6,  9],
     [7, 10],
     [8, 11]]
])

Is there a way to do this with just reshape and similar functions? The only way I can think of is using loops and copying into a new tensor.

pytorch reshape python
2 Answers

You need to use permute as well as reshape:

x.reshape(2, 2, 3).permute(0, 2, 1)
Out[*]:
tensor([[[ 0,  3],
         [ 1,  4],
         [ 2,  5]],

        [[ 6,  9],
         [ 7, 10],
         [ 8, 11]]])

First, you split the vectors into 2 x.reshape(2,2,3) placing the extra dimension in the middle. Then using permute you change the order of dimensions to be as you expected.

user_17144100

You can also use torch.split and torch.stack like

torch.stack(x.split(2), dim=2)      # or torch.stack(x.T.split(2, dim=1))
tensor([[[ 0,  3],
         [ 1,  4],
         [ 2,  5]],

        [[ 6,  9],
         [ 7, 10],
         [ 8, 11]]])
user_60756990
Related