torch.squeeze

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

返回一个移除了 input 中所有指定维度的尺寸为 1 的张量。

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

当给定 dim 时，仅在给定的维度（s）上执行挤压操作。如果输入形状为： $(A \times 1 \times B)$ ， squeeze(input, 0) 将保持张量不变，但 squeeze(input, 1) 将将张量挤压成形状 $(A \times B)$ 。

注意

返回的张量与输入张量共享存储，因此更改一个的内容将更改另一个的内容。

警告

如果张量具有大小为 1 的批处理维度，则 squeeze(input) 也会删除批处理维度，这可能导致意外的错误。请仅指定要挤压的维度。

参数:

• input (Tensor) – 输入张量。

• dim (int 或 int 的元组，可选) -

  如果提供，则输入将被压缩

  仅在指定的维度内。

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