squeeze的中文意思是"挤压 ",顾名思义,该函数的作用是压缩维度
python
squeeze(input, dim=None) -> Tensor
input一个高维张量,如果各个维度中存在大小为1的维度,squeeze才起作用,下面举例说明
python
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张量无效。