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