Python PyTorch squeeze用法及代码示例

本文简要介绍python语言中 torch.squeeze 的用法。

用法:

torch.squeeze(input, dim=None, *, out=None) → Tensor

参数：

• input(Tensor) -输入张量。

• dim(int,可选的) -如果给定，输入将只在这个维度上被压缩

关键字参数：

out(Tensor,可选的) -输出张量。

返回一个张量，其中删除了大小为 1 的 input 的所有维度。

例如，如果 input 的形状为：(A \times 1 \times B \times C \times 1 \times D)，那么 out 张量的形状为：(A \times B \times C \times D)。

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

注意

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

警告

如果张量具有大小为 1 的批量维度，那么 squeeze(input) 也会删除批量维度，这可能会导致意外错误。

例子：

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