torch中squeeze() 和unsqueeze() 的作用

torch.squeeze

torch.squeeze(input, dim=None, out=None)

将输入张量形状中的1 去除并返回。 如果输入是形如(A×1×B×1×C×1×D)，那么输出形状就为： (A×B×C×D)

当给定dim时，那么挤压操作只在给定维度上。例如，输入形状为: (A×1×B), squeeze(input, 0) 将会保持张量不变，只有用 squeeze(input, 1)，形状会变成 (A×B)。

注意： 返回张量与输入张量共享内存，所以改变其中一个的内容会改变另一个。

参数:

• input (Tensor) – 输入张量
• dim (int, optional) – 如果给定，则input只会在给定维度挤压
• out (Tensor, optional) – 输出张量

例子：

>>> x = torch.zeros(2,1,2,1,2)
>>> x.size()
(2L, 1L, 2L, 1L, 2L)
>>> y = torch.squeeze(x)
>>> y.size()
(2L, 2L, 2L)
>>> y = torch.squeeze(x, 0)
>>> y.size()
(2L, 1L, 2L, 1L, 2L)
>>> y = torch.squeeze(x, 1)
>>> y.size()
(2L, 2L, 1L, 2L)

 

torch.unsqueeze

torch.unsqueeze(input, dim, out=None)

返回一个新的张量，对输入的制定位置插入维度 1

注意： 返回张量与输入张量共享内存，所以改变其中一个的内容会改变另一个。

如果dim为负，则将会被转化dim+input.dim()+1

参数:

• tensor (Tensor) – 输入张量
• dim (int) – 插入维度的索引
• out (Tensor, optional) – 结果张量

>>> x = torch.Tensor([1, 2, 3, 4])
>>> torch.unsqueeze(x, 0)
 1  2  3  4
[torch.FloatTensor of size 1x4]
>>> torch.unsqueeze(x, 1)
 1
 2
 3
 4
[torch.FloatTensor of size 4x1]