PyTorch Tensor 形状操作

PyTorch Tensor 形状操作

我们在搭 PyTorch 网络的时候,数据要在层与层之间传来传去,而每一层需要的张量形状往往不一样。卷积层输出的是带通道维的四维张量,全连接层却要求一维的特征向量,池化之后还要把矩阵拉平才能接上分类头。形状一旦对不上,前向传播第一步就会报错。形状操作写不顺,调试网络的时间大半都耗在搬维度上。

这一节我们把最常用的四类形状操作讲透,分别是 reshape 重塑形状,transpose 和 permute 交换维度,view 和 contiguous 视图与内存,squeeze 和 unsqueeze 压缩扩展维度。每个操作我们都拿具体的张量输入走一遍调用和输出,再用表格对照前后的 shape。在开始之前,大家只需要知道张量是什么、shape 属性怎么读,以及 dim 是从 0 开始编号的。

1. 形状重塑

我们先从最日常的需求说起。数据一个元素都不删不增,只是把它们重新排列成想要的形状,reshape 做的就是这件事,它的前提是新旧形状的元素总数相等。例如:原来 2 行 3 列一共 6 个数,重排成 1 行 6 列还是 6 个数,数据本身没有动,只是摆放方式变了。

import torch

data = torch.tensor([[10, 20, 30], [40, 50, 60]])

# shape 属性和 size() 方法都能看形状
print(data.shape, data.size(0), data.size(1))

# 元素总数仍是 6,重排成 (1, 6)
print(data.reshape(1, 6).shape)

# 某个维度写 -1,PyTorch 按元素总数自动推算
print(data.reshape(2, -1).shape)
print(data.reshape(-1).shape)

运行结果如下:

torch.Size([2, 3]) 2 3
torch.Size([1, 6])
torch.Size([2, 3])
torch.Size([6])

我们把这几次调用的形状变化列在一起看:

调用输入 shape输出 shape
reshape(1, 6)(2, 3)(1, 6)
reshape(2, -1)(2, 3)(2, 3)
reshape(-1)(2, 3)(6,)

-1 是写代码时的省事写法,我们只管指定其他维度,剩下那一维交给 PyTorch 自己推算。这里要注意的是,reshape 是否复制内存是不确定的,内存刚好能复用就返回视图,复用不了才拷贝一份。所以写代码时,既不要假设它一定共享内存,也不要假设它一定拷贝。

2. 维度交换

reshape 只是把数据重新打包,有时候我们需要的是真正把两个轴对调。例如:张量存的是 (N, 通道, 高, 宽),可画图库要求的却是 (N, 高, 宽, 通道),这时候就必须交换维度,而不是 reshape。其中 transpose 一次交换两个维度,permute 一次性给出完整的新维度顺序。

import torch

data = torch.arange(24).reshape(2, 3, 4)
print(data.shape)                          # torch.Size([2, 3, 4])

# transpose 一次只交换第 1、2 维
print(torch.transpose(data, 1, 2).shape)    # torch.Size([2, 4, 3])

# permute 直接给出新顺序 [1, 2, 0]
print(torch.permute(data, [1, 2, 0]).shape) # torch.Size([3, 4, 2])

我们看这两次调用的形状变化:

调用含义输出 shape
transpose(data, 1, 2)对调第 1、2 维(2, 4, 3)
permute(data, [1, 2, 0])按 1、2、0 重排全部维度(3, 4, 2)

为什么交换维度这么快?因为 transpose 和 permute 并没有真的搬动内存里的数,只是改了这张张量的步长记录,告诉 PyTorch 现在按另一个顺序去读内存。这当然省内存,但也带来一个副作用,交换过维度的张量在内存里不再是连续排布的,这一点下一节马上就会用到。

3. 视图与内存

说到内存连续,就轮到 view 登场了。view 也是改形状,但它比 reshape 挑剔,只接受在整块内存里连续存放的张量。刚建出来的张量默认是连续的,经过 transpose、permute 之后就不再连续,这时候直接调用 view 就会报错。

view 和 reshape 到底有什么区别?

区别在于 view 要求张量在内存里必须是连续的,而 reshape 没有这个限制。我们先看下面的例子。

import torch

data = torch.tensor([[10, 20, 30], [40, 50, 60]])

# 连续张量可以直接 view
print(data.view(3, 2).shape)     # torch.Size([3, 2])
print(data.is_contiguous())      # True

# transpose 之后变成非连续
t = torch.transpose(data, 0, 1)
print(t.is_contiguous())         # False
# t.view(3, 2) 这一行会抛 RuntimeError

# 先 contiguous() 把数据拷成连续排布,再 view
print(t.contiguous().view(3, 2).shape)  # torch.Size([3, 2])

程序输出依次是:

torch.Size([3, 2])
True
False
torch.Size([3, 2])

is_contiguous 返回 True,表示数据在内存里首尾相接排好了,返回 False 就说明中间有跳步。非连续的时候想再 view,要先调一次 contiguous,它会真的把数据重新拷成连续布局,拷贝花的时间和张量大小成正比。日常写代码其实可以直接用 reshape,它内部会自动判断,连续就返回视图,不连续就自动拷贝,省得我们手动处理。view 则留给那些明确知道数据连续、又想确保不触发拷贝的场合。

4. 维度压缩

还有一种常见的麻烦,是张量里混着大小为 1 的多余维度。例如:卷积输出多挂了一个批维,或者某个中间结果的形状是 (1, 3, 1, 5),而真正要算的只是里面那个 3 乘 5 的矩阵。squeeze 负责把大小为 1 的维度挤掉,unsqueeze 则反过来,在指定位置塞一个大小为 1 的维度进去。

import torch

data = torch.arange(15).reshape(1, 3, 1, 5)
print(data.shape)                 # torch.Size([1, 3, 1, 5])

# 不带参数,挤掉所有大小为 1 的维
print(data.squeeze().shape)       # torch.Size([3, 5])

# 指定位置挤,那维必须恰好是 1 才生效
print(data.squeeze(0).shape)      # torch.Size([3, 1, 5])
print(data.squeeze(2).shape)      # torch.Size([1, 3, 5])
print(data.squeeze(1).shape)      # torch.Size([1, 3, 1, 5]),第 1 维是 3,不动

# 在最后加一个大小为 1 的维
print(data.unsqueeze(-1).shape)   # torch.Size([1, 3, 1, 5, 1])

我们把各次调用的形状对照列在下表:

调用输出 shape说明
squeeze()(3, 5)挤掉所有大小为 1 的维
squeeze(0)(3, 1, 5)第 0 维是 1,挤掉
squeeze(2)(1, 3, 5)第 2 维是 1,挤掉
squeeze(1)(1, 3, 1, 5)第 1 维是 3,保持不变
unsqueeze(-1)(1, 3, 1, 5, 1)末尾加一维

关键在于 squeeze(dim) 是一个条件操作。我们指定的那一维如果确实是 1,就挤掉,不是 1 就原样返回,既不会报错也不会乱删。配合广播做矩阵乘法时,我们经常用 unsqueeze 手动补一个 1,让两个张量的维度对齐。

5. 常见误区

第一类坑,是把形状操作当成复制。reshape、view、transpose、squeeze、unsqueeze 默认都和原张量共享同一块内存,改了新张量里的数,原张量会跟着变。需要一份独立数据时,要先调 .clone() 再改,否则调试时会出现意想不到的联动。

第二类坑,是纠结 view 和 reshape 该用哪个。记住一条就够,不确定数据连不连续就用 reshape,它最省心。只有当我们明确要在连续张量上做零拷贝视图、又想让编译器帮忙检查这一点时,才用 view。

第三类坑,是 reshape 前后元素总数对不上。(2, 3) 一共 6 个数,reshape(1, 6) 可以,reshape(1, 5) 就会直接报错。报错信息里会写清楚期望的元素个数,对不上时先拿新旧 shape 各乘一遍核对。

形状操作说到底就是三件事,改总布局用 reshape,对调坐标轴用 transpose 或 permute,挤掉或补上大小为 1 的维用 squeeze 和 unsqueeze。transpose 不搬数据只改步长,代价是让张量失去内存连续性,这时再用 view 就要先 contiguous。日常写代码优先用 reshape 和 squeeze,把内存细节交给 PyTorch,等调试报错时再回头查 is_contiguous,网络层之间的形状衔接就顺了。