Pytorch squeeze()的简单用法

squeeze的作用就是对tensor变量进行维度压缩，去除维数为1的的维度。
☆例如，一个3×2×1×2×1的tensor，squeeze()之后便成了3×2×2。存储的数据并没有发生变化，但是去除了“多余”的维度信息：

import torch

x = torch.rand(3, 2, 1, 2, 1)
print(x.size())
x = x.squeeze()
print(x.size())

输出：

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

☆如果原tensor没有维度为1，则不会进行任何操作：

import torch

x = torch.rand(2, 2)
print(x)
x = x.squeeze()
print(x)

输出：

tensor([[0.6088, 0.5360],
        [0.4375, 0.0350]])
tensor([[0.6088, 0.5360],
        [0.4375, 0.0350]])

☆此外，也可以指定尝试将哪一个维度进行压缩。如果被指定的维度其维数为1，则压缩，反之不对该维度操作：

import torch

x = torch.rand(3, 2, 1, 2, 1)
print(x.size())
x = x.squeeze(0)
print(x.size())
x = x.squeeze(2)
print(x.size())

输出：

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

☆torch中的squeeze方法与上面提到的Tensor对象中的方法作用相同，语法有细微区别：

import torch
x = torch.rand(3, 2, 1, 2, 1)
x = torch.squeeze(x) # 同x.squeeze()
print(x.size())
x = torch.squeeze(x, 2) # 同x.squeeze(2)
print(x.size())