torch.squeeze¶
- torch.squeeze(input: Tensor, dim: Optional[Union[int, List[int]]]) Tensor ¶
返回一个新张量,其中移除了
input
中所有大小为 1 的指定维度。例如,如果 input 的形状是:,则 input.squeeze() 的形状将是:。
当指定了
dim
时,squeeze 操作仅在指定的维度上执行。如果 input 的形状是:,则squeeze(input, 0)
会使张量保持不变,而squeeze(input, 1)
会将张量压缩为形状 。注意
返回的张量与输入张量共享底层存储,因此改变其中一个的内容会改变另一个的内容。
警告
如果张量的批处理维度大小为 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]) >>> y = torch.squeeze(x, (1, 2, 3)) torch.Size([2, 2, 2])