【pytorch算子】torch.squeeze 和 torch.unsqueeze 的用法

torch.squeeze() 用法：

对 Tensor 进行降维，去掉维数为 1 的的维度。

a.squeeze() 是将 a 中所有为 1 的维度删掉，不为 1 的维度没有影响。
a.squeeze(n) 是将 a 中第 n 个维度中维数为 1 的维度删掉，即若 dim(n) = 1，则删掉第 n 维；若 dim(n) != 1，则保留第 n 维，不做任何操作。

torch.unsqueeze() 用法：

对 Tensor 进行扩维，扩大的维数为 1 。

a.squeeze(n) 是将 a 中第 n 个维度的维数改为 1 ，其他维度依次增加 1，即扩充第 n 维。

代码示例

a=torch.tensor([
    [[1,2,3,5]],
    [[2,3,6,5]],
    [[5,9,9,6]]])      # a.shape=torch.Size([3, 1, 4])
b=a.squeeze()    # 去掉维度为1的维度， b.shape=torch.Size([3, 4])
c=b.unsqueeze(0)   # 扩充第0维， b.shape=torch.Size([1, 3, 4])
c=b.unsqueeze(1)   # 扩充第1维， b.shape=torch.Size([3, 1, 4])
c=b.unsqueeze(2)   # 扩充第2维， b.shape=torch.Size([3, 4, 1])


d=torch.tensor([[[1,2],[2,3]],
                [[5,6],[5,6]],
                [[7,5],[7,5]]])    # d.shape=torch.Size([3, 2, 2])
e=d.squeeze()    # d中没有维数为1的维度，因此shape不变，e.shape=torch.Size([3, 2, 2])
f=d.squeeze(1)   # d中第1维的维数不为1，因此无法降维，shape不变，e.shape=torch.Size([3, 2, 2])