pytorch中squeeze函数用法

squeeze的中文意思是“挤压”，顾名思义，该函数的作用是压缩维度

squeeze(input, dim=None) -> Tensor

input一个高维张量，如果各个维度中存在大小为1的维度，squeeze才起作用，下面举例说明

x = torch.arange(6).reshape(2,1,3)
 
# tensor([[[0, 1, 2]],
#         [[3, 4, 5]]])  shape=(2,1,3)
 
x = x.squeeze()
 
# tensor([[0, 1, 2],
#        [3, 4, 5]])    shape=(2,3)

如果指定dim，当dim=1时候，效果和上面相同，如果dim为其他，那么x的维度不变。

当然，如果各个维度中不存在大小为1的维度，那么squeeze函数对x张量无效。