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