torch.squeeze

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

回傳一個張量，其中已移除 input 中所有大小為 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) 會將張量壓縮為 $(A \times B)$ 形狀。

注意

返回的張量與輸入張量共享儲存空間，因此更改其中一個張量的內容將更改另一個張量的內容。

警告

如果張量具有大小為 1 的批次維度，則 squeeze(input) 也會移除批次維度，這可能會導致意外的錯誤。 考慮僅指定您希望壓縮的維度。

參數

• input (Tensor) – 輸入張量。

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

  如果給定，輸入將被壓縮

  僅在指定的維度中。

  變更於版本 2.0：dim 現在接受維度的 tuple。

範例

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