快捷方式

torch.squeeze

torch.squeeze(input: Tensor, dim: Optional[Union[int, List[int]]]) Tensor

返回一个新张量,其中移除了 input 中所有大小为 1 的指定维度。

例如,如果 input 的形状是:(A×1×B×C×1×D)(A \times 1 \times B \times C \times 1 \times D),则 input.squeeze() 的形状将是:(A×B×C×D)(A \times B \times C \times D)

当指定了 dim 时,squeeze 操作仅在指定的维度上执行。如果 input 的形状是:(A×1×B)(A \times 1 \times B),则 squeeze(input, 0) 会使张量保持不变,而 squeeze(input, 1) 会将张量压缩为形状 (A×B)(A \times B)

注意

返回的张量与输入张量共享底层存储,因此改变其中一个的内容会改变另一个的内容。

警告

如果张量的批处理维度大小为 1,则 squeeze(input) 也会移除该批处理维度,这可能导致意外错误。请考虑仅指定您希望压缩的维度。

参数
  • input (Tensor) – 输入张量。

  • dim (inttuple of ints, 可选) –

    如果给定,输入张量仅在指定的维度上

    被压缩。

    版本 2.0 中有更改: dim 现在接受维度元组。

示例

>>> 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])
>>> y = torch.squeeze(x, (1, 2, 3))
torch.Size([2, 2, 2])

文档

查阅 PyTorch 的完整开发者文档

查看文档

教程

获取针对初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得问题解答

查看资源