PyTorch Tensor 拼接操作

PyTorch Tensor 拼接操作

搭网络的时候,我们常常要把两路张量合到一起。残差网络里,主路径卷积后的特征要和捷径分支的特征加回来,注意力机制里多个头的输出要拼回完整维度,多模态模型要把图像特征和文本特征按通道接起来。合张量是个高频操作,PyTorch 里就两个主力函数:cat 和 stack,把它俩搞混是新手最常踩的坑。

这一节我们用具体张量把两个函数各走一遍,看清楚沿哪个维度合、形状怎么变,最后用一张表说清什么时候用 cat、什么时候用 stack。读之前要先知道张量的 shape 怎么读,dim 是从 0 开始编号的。

1. 按维拼接

我们先看 torch.cat,它的作用是沿着一个已经存在的维度,把若干张量首尾接起来。要求是除了被拼接的那一维,其他维度的形状必须完全一致。两个 (2, 3, 2) 的张量按第 0 维接,就变成 (4, 3, 2),按第 2 维接就变成 (2, 3, 4),接完之后总维度数不变。

import torch

a = torch.arange(12).reshape(2, 3, 2)
b = torch.arange(100, 112).reshape(2, 3, 2)
print(a.shape, b.shape)            # torch.Size([2, 3, 2])

# 沿第 0 维接,批次数翻倍
print(torch.cat([a, b], dim=0).shape)   # torch.Size([4, 3, 2])

# 沿第 1 维接
print(torch.cat([a, b], dim=1).shape)   # torch.Size([2, 6, 2])

# 沿第 2 维接
print(torch.cat([a, b], dim=2).shape)   # torch.Size([2, 3, 4])

三种拼接方向的形状变化,我们列在一起看:

拼接维度结果 shape哪一维变长了
dim=0(4, 3, 2)第 0 维 2 变 4
dim=1(2, 6, 2)第 1 维 3 变 6
dim=2(2, 3, 4)第 2 维 2 变 4

直观上理解,dim=0 是把两块样本上下摞起来,dim=2 是把特征在最后一截加宽。残差网络里把主路径和捷径特征按通道维 dim=1 拼起来,用的就是 cat。要记住,拼接只是把两块数据在指定轴上首尾相接,不会凭空多出一个新维度。

2. 堆叠升维

stack 和 cat 名字像,做的事却不一样。torch.stack 会先在指定位置插入一个全新的维度,再把各张量沿着这个新维度放进去。它要求参与堆叠的张量形状完全相同,因为每个张量都要被当成新维度上的一个切片。两个 (2, 3) 的矩阵,stack 完一定变成三维张量。

import torch

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
y = torch.tensor([[7, 8, 9], [10, 11, 12]])

# 在最前面插入新维,x 和 y 各占一层
print(torch.stack([x, y], dim=0).shape)  # torch.Size([2, 2, 3])

# 在第 1 维插入新维
print(torch.stack([x, y], dim=1).shape)  # torch.Size([2, 2, 3])

# 在最后插入新维
print(torch.stack([x, y], dim=2).shape)  # torch.Size([2, 3, 2])

我们把三种堆叠的实际数值打出来看,位置上的差别就清楚了。dim=0 堆叠时,x 整块在第一层,y 整块在第二层:

tensor([[[ 1,  2,  3],
         [ 4,  5,  6]],

        [[ 7,  8,  9],
         [10, 11, 12]]])

dim=2 堆叠时,则是把相同位置的元素两两配对,x 的 1 和 y 的 7 并到一个位置上:

tensor([[[ 1,  7],
         [ 2,  8],
         [ 3,  9]],

        [[ 4, 10],
         [ 5, 11],
         [ 6, 12]]])

收集一批同形状的样本,就是 stack 的典型场景。每张图片本身是 (通道, 高, 宽),读进来几张之后用 stack 沿 dim=0 堆成 (批次数, 通道, 高, 宽),这个额外的批次维就是 stack 凭空加出来的。

3. 对比与选择

两个函数的差别一句话就能说清:cat 在已有维度上延长,stack 先新建一个维度再往里放。把关键性质对照起来,选的时候就不纠结了:

对比项torch.cattorch.stack
结果维度数和输入相同比输入多一维
形状要求除拼接维外其余维一致所有维度完全一致
(2,3) 加 (2,3)(4,3) 或 (2,6)(2,2,3) 或 (2,3,2)
典型用途特征拼接、通道合并组批次、捆样本

判断时问自己一个问题:我是想把数据接长一点,还是想给这批数据加一个批次号。想接长就用 cat,想编号就用 stack。

4. 常见误区

cat 最常见的报错,是被拼接维度之外的形状对不上。两个张量要在 dim=1 接,那它们的第 0 维和第 2 维必须相等,差一个数 PyTorch 就抛 RuntimeError。报错信息会把两边的 shape 打出来,对照着找差在哪一维就行。

stack 的坑在于形状必须严丝合缝。它比 cat 更苛刻,所有维度都要一致,因为每个张量都要被当成新维度上对称的一层。形状不齐时,先 reshape 或 pad 到一致,再 stack。

第三个误区是用 stack 替代 cat 去拼特征。特征融合要的是通道变长、维度数不变,用 stack 凭空多加一维,会把后面的全连接层形状搞乱。记住:该接长就接长,该编号才编号。

小结一下。cat 沿已有维度把张量接长,维度数不变,除拼接维外形状要一致,适合残差分支、通道融合这类特征拼接。stack 先插一个新维度,再把同形状张量放进去,维度数加一,适合把一批样本捆成一个批次。只要记住维度数变不变这一条,两个函数就不会用反。