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 torch

x1 = 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