Pytorch：squeeze()和unsqueeze()函数

squeeze(): 维度压缩，去掉维数为1的的维度

torch.squeeze(input, dim=None, out=None)
将输入张量形状中的1 去除并返回。 如果输入是形如(2×1×2×1×2×1×2)，那么输出形状就为： (2×2×2×2)。
a.squeeze(N)就是去掉a中指定的维数为一的维度。还有一种形式就是b=torch.squeeze(a，N) a中去掉指定的定的维数为一的维度。

参数:
input (Tensor) – 输入张量
dim (int, optional) – 如果给定，则input只会在给定维度挤压
out (Tensor, optional) – 输出张量

a = torch.zeros(3,1,2,1,2)
a.size()
# torch.Size([3, 1, 2, 1, 2])
y = torch.squeeze(a)
y.size()
# torch.Size([3, 2, 2])
y = torch.squeeze(a, 0)
y.size()
# torch.Size([3, 1, 2, 1, 2])
y = torch.squeeze(a, 1)
y.size()
# torch.Size([3, 2, 1, 2])
y = a.squeeze(1)
y.size()
# torch.Size([3, 2, 1, 2])

---

unsqueeze():扩充数据维度，在指定位置N加上维数为1的维度

torch.unsqueeze(input, dim, out=None)
返回一个新的张量，对输入的指定位置插入维度 1。
注意：如果dim为负，则将会被转化dim+input.dim()+1

参数:
tensor (Tensor) – 输入张量
dim (int) – 插入维度的索引
out (Tensor, optional) – 结果张量

x = torch.Tensor([1, 2, 3, 4])
x.size()
# torch.Size([4])
y = torch.unsqueeze(x, 0)
y.size()
# torch.Size([1, 4])
y = torch.unsqueeze(x, 1)
y.size()
# torch.Size([4, 1])
y = x.unsqueeze(1)
y.size()
# torch.Size([4, 1])