pytorch学习——torch.cat和torch.stack的区别

合并tensors

torch.cat 沿着特定维数连接一系列张量。
torch.stack 沿新维度连接一系列张量。

torch.cat

在给定维度中连接给定的 seq 个张量序列。
所有张量必须具有相同的形状（连接维度除外）或为空。

torch.cat(tensors, dim=0, *, out=None) → Tensor

参数

tensors（张量序列）：任何相同类型的张量序列。提供的非空张量必须具有相同的形状。在给定维度上对输入的张量序列进行连接操作。
dim (int) ：张量连接的维度，

torch.stack

沿新维度连接一系列张量。（维度叠加）
所有张量都需要具有相同的大小。

torch.stack(tensors, dim=0, *, out=None) → Tensor

参数

tensors（张量序列）：要连接的张量序列
dim (int) ：要插入的维度。必须介于 0 和串联张量的维数之间（含）

示例

沿第0维操作：

import torchx1 = torch.tensor([[1,2,3], [4,5,6]])# x1.shape = tensor.size([2,3])
x2 = torch.tensor([[7,8,9], [10,11,12]])# x2.shape = tensor.size([2,3])
print(x1.shape)
print('沿第0维进行操作：')
y1 = torch.cat([x1, x2], dim=0)
y2 = torch.stack([x1, x2], dim=0)
print('cat, y1:', y1.shape,'\n',y1)
print('stack, y2:', y2.shape,'\n',y2)

输出：

沿第0维进行操作：
cat, y1: torch.Size([4, 3]) tensor([[ 1,  2,  3],[ 4,  5,  6],[ 7,  8,  9],[10, 11, 12]])
stack, y2: torch.Size([2, 2, 3]) tensor([[[ 1,  2,  3],[ 4,  5,  6]],[[ 7,  8,  9],[10, 11, 12]]])

从y1的输出可以看到，cat在第0维将x1和x2元素进行续接，即输出为[x1[0], x1[1], x2[0], x2[1]]， shape由[2, 3]变为[4,3]。

从y2的输出可以看到，stack直接将x1和x2的第0维进行叠加，即输出为[x1, x2]，shape由[2,3]变为[2, 2, 3]。

沿第1维操作：

print('沿第1维进行操作：')
y1 = torch.cat(x, dim=1)
y2 = torch.stack(x, dim=1)
print('cat, y1:', y1.shape,'\n',y1)
print('stack, y2:', y2.shape,'\n',y2)

输出：

沿第1维进行操作：
cat, y1: torch.Size([2, 6]) tensor([[ 1,  2,  3,  7,  8,  9],[ 4,  5,  6, 10, 11, 12]])
stack, y2: torch.Size([2, 2, 3]) tensor([[[ 1,  2,  3],[ 7,  8,  9]],[[ 4,  5,  6],[10, 11, 12]]])

从y1的输出可以看到，cat将x1和x2相对应的第1维的元素进行续接， shape由[2,3]变为[2, 6]。

从y2的输出可以看到，stack直接将x1和x2相对应的第1维的元素进行叠加，即输出为[[x1[0], x2[0]], [x1[1], x2[1]]，shape由[2,3]变为[2, 2, 3]。

沿第2维操作：

输出

y1 = torch.cat(x, dim=2)
print('cat, y1:', y1.shape,'\n',y1)Traceback (most recent call last):File "/Users/gyuer/Desktop/test.py", line 8, in <module>y1 = torch.cat(x, dim=2)
IndexError: Dimension out of range (expected to be in range of [-2, 1], but got 2)

y2 = torch.stack(x, dim=2)
print('stack, y2:', y2.shape,'\n',y2)stack, y2: torch.Size([2, 3, 2]) tensor([[[ 1,  7],[ 2,  8],[ 3,  9]],[[ 4, 10],[ 5, 11],[ 6, 12]]])

从以上结果可以看出，torch.stack(x, dim=2)是将x1[i][j]和x2[i][j]堆叠在一起的。如x1[0][0]=1和x2[0][0]=7堆叠在一起，得到[1, 7]。
stack的参数dim要插入的维度必须介于 0 和串联张量的维数之间

以上总结借鉴了官网的英文解释和https://blog.csdn.net/weixin_42920104/article/details/105833691