torch.squeeze

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

返回一个张量，其中所有大小为 1 的指定维度都已被移除。

例如，如果 input 的形状为： $(A \times 1 \times B \times C \times 1 \times D)$，那么 input.squeeze() 的形状将是： $(A \times B \times C \times D)$。

当指定 dim 时，仅在给定的维度上执行 squeeze 操作。如果 input 的形状为： $(A \times 1 \times B)$，则 squeeze(input, 0) 将保持张量不变，但 squeeze(input, 1) 将把张量 squeeze 到形状 $(A \times B)$。

注意

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

警告

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

参数

• input (Tensor) – 输入张量。

• dim (int 或 tuple of ints, optional) –

  如果给定，则仅在指定的维度上对输入执行 squeeze 操作。

  仅在指定的维度上。

  版本 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])