torch.stack()用法

torch.stack((tensor1, tensor2,…), dim=x, out=None)
作用：于将tensor进行按照指定维度进行堆叠。

一维tensor

w = torch.arange(4)
h = torch.tensor([0.6, 0.7, 0.8, 0.2])
torch.stack([-w, -h, w, h], dim=0)
>>>
tensor([[ 0.00, -1.00, -2.00, -3.00],
        [-0.60, -0.70, -0.80, -0.20],
        [ 0.00,  1.00,  2.00,  3.00],
        [ 0.60,  0.70,  0.80,  0.20]])
# dim = 1
w = torch.arange(4)
h = torch.tensor([0.6, 0.7, 0.8, 0.2])
torch.stack((w, h), dim=1)
>>>
tensor([[0.00, 0.60],
        [1.00, 0.70],
        [2.00, 0.80],
        [3.00, 0.20]])

二维tensor

dim = 1时，将tensor1[0]和tensor2[0]进行堆叠；
dim = 2时，将tensor1[0][0]和tensor2[0][0]进行堆叠；

w = torch.arange(12).reshape(3, 4)
h = -torch.arange(12).reshape(3, 4)
c = torch.stack((w, h), dim=0)
>>>
tensor([[[  0,   1,   2,   3],
         [  4,   5,   6,   7],
         [  8,   9,  10,  11]],

        [[  0,  -1,  -2,  -3],
         [ -4,  -5,  -6,  -7],
         [ -8,  -9, -10, -11]]])
>>> c.shape
>torch.Size([2, 3, 4])
-----------------------------------------------
b = torch.stack((w, h), dim=1)
>>>
tensor([[[  0,   1,   2,   3],
         [  0,  -1,  -2,  -3]],

        [[  4,   5,   6,   7],
         [ -4,  -5,  -6,  -7]],

        [[  8,   9,  10,  11],
         [ -8,  -9, -10, -11]]])
>>>b.shape
>torch.Size([3, 2, 4])
-----------------------------------------------
d = torch.stack((w, h), dim=2)
print(d)
print(d.shape)
>>>
tensor([[[  0,   0],
         [  1,  -1],
         [  2,  -2],
         [  3,  -3]],

        [[  4,  -4],
         [  5,  -5],
         [  6,  -6],
         [  7,  -7]],

        [[  8,  -8],
         [  9,  -9],
         [ 10, -10],
         [ 11, -11]]])
torch.Size([3, 4, 2])