【PyTorch】squeeze函数详解

#【PyTorch】squeeze函数详解

函数原型

torch.squeeze(input, dim=None, *, out=None) → Tensor

函数详解

删除一个张量中所有维数为1的维度

例如，一个维度为 $(A\times 1\times B\times C\times 1\times D)$的张量调用squeeze方法后维度变为$(A\times B\times C\times D)$

当给定dim参数后，squeeze操作只作用于给定维度。如果输入input的形状为$(A\times 1\times B)$, squeeze(input, 0)不改变这个张量, 但是squeeze(input, 1) 将把这个张量的形状变为$(A\times B)$.

举例

>>> x = torch.zeros(2, 1, 2, 1, 2)
>>> x.size()
torch.Size([2, 1, 2, 1, 2])
>>> y = torch.squeeze(x)
>>> y.size()
torch.Size([2, 2, 2])
>>> y = torch.squeeze(x, 0)
>>> y.size()
torch.Size([2, 1, 2, 1, 2])
>>> y = torch.squeeze(x, 1)
>>> y.size()
torch.Size([2, 2, 1, 2])

————文章内容翻译自[PyTorch官方文档](https://pytorch.org/docs/stable/generated/torch.squeeze.html?highlight=squeeze#torch.squeeze)