pytorch: 给tensor删除或者添加维度为1的维度（squeeze和unsqueeze）

删除tensor中维度为1的维度

在pytorch中，用torch.squeeze()函数或者tensor的自身成员函数squeeze()去除维度为1的维度。

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])

给tensor添加维度为1的维度

在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])