[pytoch foundation] torch.stack() function analysis

1 Function function

   Use stack Two messages can be retained ：[1. Sequence ] and [2. Tensor matrix ] Information , Belong to Expand and splice Function of ; It can be considered that one matrix is compressed into a matrix according to the time series . It often appears in natural language processing （NLP） Image convolution neural network (CV) in .
   Official explanation ： Connect the sequence of input tensors along a new dimension . All tensors in the sequence should be of the same shape .
   In simple terms ： The multiple 2 The tensor of the dimension comes together into a 3 D tensor ; Multiple 3 Dimension makes up a 4 D tensor … And so on , That is, adding new dimensions for stacking .

2 Example

import torch
import numpy as np
#  establish 3*3 Matrix ,a、b
a=np.array([[1,2,3],[4,5,6],[7,8,9]])
b=np.array([[10,20,30],[40,50,60],[70,80,90]])
#  Transform the matrix into Tensor
a = torch.from_numpy(a)
b = torch.from_numpy(b)

tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]], dtype=torch.int32)
tensor([[10, 20, 30],
        [40, 50, 60],
        [70, 80, 90]], dtype=torch.int32)

2.1 dim=0

d = torch.stack((a, b), dim=0)
print(d)
print(d.size())
tensor([[[ 1,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 20, 30],
         [40, 50, 60],
         [70, 80, 90]]], dtype=torch.int32)
torch.Size([2, 3, 3])

   When dim = 0, Every original matrix has also become a dimension , See a matrix as a whole . There are several matrices , The new dimension is , The matrix is the dimension .

2.2 dim=1

d = torch.stack((a, b), dim=1)
print(d)
print(d.size())
tensor([[[ 1,  2,  3],
         [10, 20, 30]],

        [[ 4,  5,  6],
         [40, 50, 60]],

        [[ 7,  8,  9],
         [70, 80, 90]]], dtype=torch.int32)
torch.Size([3, 2, 3])

   Form the first row of each matrix into a first dimensional matrix , Go down in turn , The... Of each matrix n Line composition No n D matrix .size=(n,i,y)

2.3 dim=2

d = torch.stack((a, b), dim=2)
print(d)
print(d.size())
tensor([[[ 1, 10],
         [ 2, 20],
         [ 3, 30]],

        [[ 4, 40],
         [ 5, 50],
         [ 6, 60]],

        [[ 7, 70],
         [ 8, 80],
         [ 9, 90]]], dtype=torch.int32)
torch.Size([3, 3, 2])

3 reference

[1]【Pytorch】torch.stack() Use
[2] After reading the second understand torch.stack()