当前位置：首页 > news >正文

Pytorch中stack()方法的总结及理解

news 2025/9/25 16:26:08

Torch.stack()
1. 概念
在一个新的维度上连接一个张量序列

2. 参数
tensors （sequence）需要连接的张量序列
dim （int）在第dim个维度上连接
注意输入的张量shape要完全一致

3. 举例
3.1 四个shape为[3, 3]的张量
以下面这4个张量，每个张量shape为[3, 3]。

a = torch.Tensor([[1,2,3],[4,5,6],[7,8,9]])
b = torch.Tensor([[10,20,30],[40,50,60],[70,80,90]])
c = torch.Tensor([[100,200,300],[400,500,600],[700,800,900]])
d = torch.Tensor([[1000,2000,3000],[4000,5000,6000],[7000,8000,9000]])

3.1.1 dim=0的情况下，直接来看结果。

torch.stack([a,b,c,d], dim=0)

tensor([[[1.0000e+00, 2.0000e+00, 3.0000e+00],[4.0000e+00, 5.0000e+00, 6.0000e+00],[7.0000e+00, 8.0000e+00, 9.0000e+00]],[[1.0000e+01, 2.0000e+01, 3.0000e+01],[4.0000e+01, 5.0000e+01, 6.0000e+01],[7.0000e+01, 8.0000e+01, 9.0000e+01]],[[1.0000e+02, 2.0000e+02, 3.0000e+02],[4.0000e+02, 5.0000e+02, 6.0000e+02],[7.0000e+02, 8.0000e+02, 9.0000e+02]],[[1.0000e+03, 2.0000e+03, 3.0000e+03],[4.0000e+03, 5.0000e+03, 6.0000e+03],[7.0000e+03, 8.0000e+03, 9.0000e+03]]])

此时在第0个维度上连接，新张量的shape可以发现为[4, 3, 3]，4代表在第0个维度有4项。

观察可以得知：即初始的四个张量，即a、b、c、d四个初始张量。

可以理解为新张量的第0个维度上连接a、b、c、d。

3.1.2 dim=1的情况下

torch.stack([a,b,c,d], dim=1)

tensor([[[1.0000e+00, 2.0000e+00, 3.0000e+00],[1.0000e+01, 2.0000e+01, 3.0000e+01],[1.0000e+02, 2.0000e+02, 3.0000e+02],[1.0000e+03, 2.0000e+03, 3.0000e+03]],[[4.0000e+00, 5.0000e+00, 6.0000e+00],[4.0000e+01, 5.0000e+01, 6.0000e+01],[4.0000e+02, 5.0000e+02, 6.0000e+02],[4.0000e+03, 5.0000e+03, 6.0000e+03]],[[7.0000e+00, 8.0000e+00, 9.0000e+00],[7.0000e+01, 8.0000e+01, 9.0000e+01],[7.0000e+02, 8.0000e+02, 9.0000e+02],[7.0000e+03, 8.0000e+03, 9.0000e+03]]])

此时在第1个维度上连接，新张量的shape可以发现为[3, 4, 3]，4代表在第1个维度有4项。

观察可以得知：

新张量[0][0]为a[0]，[0][1]为b[0]，[0][2]为c[0]，[0][3]为d[0]
新张量[1][0]为a[1]，[1][1]为b[1]，[1][2]为c[1]，[1][3]为d[1]
新张量[2][0]为a[2]，[2][1]为b[2]，[2][2]为c[2]，[2][3]为d[2]

可以理解为新张量的第1个维度上连接a、b、c、d的第0个维度单位，具体地说，在新张量[i]中连接a[i]、b[i]、c[i]、d[i]，即将a[i]赋给新张量[i][0]、b[i]赋给新张量[i][1]、c[i]赋给新张量[i][2]、d[i]赋给新张量[i][3]。

3.1.2 dim=2的情况下

torch.stack([a,b,c,d], dim=2)

tensor([[[1.0000e+00, 1.0000e+01, 1.0000e+02, 1.0000e+03],[2.0000e+00, 2.0000e+01, 2.0000e+02, 2.0000e+03],[3.0000e+00, 3.0000e+01, 3.0000e+02, 3.0000e+03]],[[4.0000e+00, 4.0000e+01, 4.0000e+02, 4.0000e+03],[5.0000e+00, 5.0000e+01, 5.0000e+02, 5.0000e+03],[6.0000e+00, 6.0000e+01, 6.0000e+02, 6.0000e+03]],[[7.0000e+00, 7.0000e+01, 7.0000e+02, 7.0000e+03],[8.0000e+00, 8.0000e+01, 8.0000e+02, 8.0000e+03],[9.0000e+00, 9.0000e+01, 9.0000e+02, 9.0000e+03]]])