torch.squeeze() 和torch.unsqueeze()的用法，给tensor删除或者添加维度为1的维度（squeeze和unsqueeze）

import torch
x = torch.randn(3,1,4,1,2)
a = x.squeeze(dim=1)  # 成员函数删除第二维度
print(a.shape)
b = torch.squeeze(x, dim=1)  # torch自带函数删除第二维度
print(b.shape)
c = torch.squeeze(x, dim=3)  # 删除第三维度
print(c.shape)
d = torch.squeeze(x)  # 若不标注删除第几维度，则会删除所有为1的维度
print(d.shape)

输出：
torch.Size([3, 4, 1, 2])
torch.Size([3, 4, 1, 2])
torch.Size([3, 1, 4, 2])
torch.Size([3, 4, 2])

注意：squeeze()函数只能删除维度为1的维度，若不为1，该操作无效，但不会报错

b = torch.unsqueeze(a, dim=1)
print(b.shape)

输出为：
torch.Size([3, 1, 4, 1, 2])

 

在pytorch中， 用自带的torch.unsqueeze()和tensor的成员函数unsqueeze()可以为tensor添加维度为1的维度

e = torch.unsqueeze(x, dim=0) # 在第一维度添加维度
print(e.shape)
f = x.unsqueeze(dim=0)
print(f.shape)

输出：
torch.Size([1, 3, 1, 4, 1, 2])
torch.Size([1, 3, 1, 4, 1, 2])

 

先看torch.squeeze() 这个函数主要对数据的维度进行压缩，去掉维数为1的的维度，比如是一行或者一列这种，一个一行三列（1,3）的数去掉第一个维数为一的维度之后就变成（3）行。squeeze(a)就是将a中所有为1的维度删掉。不为1的维度没有影响。a.squeeze(N) 就是去掉a中指定的维数为一的维度。还有一种形式就是b=torch.squeeze(a，N) a中去掉指定的定的维数为一的维度。

再看torch.unsqueeze()这个函数主要是对数据维度进行扩充。给指定位置加上维数为一的维度，比如原本有个三行的数据（3），在0的位置加了一维就变成一行三列（1,3）。a.squeeze(N) 就是在a中指定位置N加上一个维数为1的维度。还有一种形式就是b=torch.squeeze(a，N) a就是在a中指定位置N加上一个维数为1的维度