pytorch 笔记之squeeze

squeeeze使用

torch.squeeze(input, dim=None, out=None) → Tensor
Returns a tensor with all the dimensions of input of size 1 removed.

简单来说，squeeze只压缩所有长度为1的维度。如果指定dim(int, optional)，则在指定的维度上压缩，且只在长度为1时起作用。

例子：

a = torch.zeros(1,2,3,4,5,1)
a.size()

# output
torch.Size([1, 2, 3, 4, 5, 1])

默认用法，去掉所有长度为1的维度：

torch.squeeze(a).size()

# output
orch.Size([2, 3, 4, 5])

指定dim：

# 压缩第一个维度（第一个维度长度是1，work）
torch.squeeze(a, dim=0).size()
# output
torch.Size([2, 3, 4, 5, 1])

# 压缩第二个维度（第二个维度长度是2，don't work）
torch.squeeze(a, dim=1).size()
# output
torch.Size([1, 2, 3, 4, 5, 1])

# 压缩第六个维度（第六个维度长度是1，work）
torch.squeeze(a, dim=5).size()
# output
torch.Size([1, 2, 3, 4, 5])