torch.unsqueeze()函数理解
torch.unsqueeze()函数理解
torch.unsqueeze(input, dim) 使用时等同于 input.unsqueeze(dim)
torch.unsqueeze()函数起到升维的作用，dim等于几表示在第几维度加一，比如原来x的size=([4])，x.unsqueeze(0)之后就变成了size=([1, 4])，而x.unsqueeze(1)之后就变成了size=([4, 1])，注意dim∈[-input.dim() - 1, input.dim() + 1]
例如
- 输入一维张量，即input.dim()=1
# 输入：
x = torch.tensor([1, 2, 3, 4]) # x.dim()=1
print(x)
print(x.shape)
y = x.unsqueeze(0)
print(y)
print(y.shape) # 此时y.dim()=2
z = x.unsqueeze(1)
print(z)
print(z.shape) # 此时z.dim()=2
# 输出：
tensor([1, 2, 3, 4])
torch.Size([4])
tensor([[1, 2, 3, 4]])
torch.Size([1, 4])
tensor([[1],
[2],
[3],
[4]])
torch.Size([4, 1])
- 输入二维张量，即input.dim()=2
# 输入：
x = torch.tensor([[1, 2, 3], [4, 5, 6]]) # x.dim()=2
print(x)
print(x.shape)
y = x.unsqueeze(0)
print(y)
print(y.shape) # 此时y.dim()=3
z = x.unsqueeze(1)
print(z)
print(z.shape) # 此时z.dim()=3
# 输出：
tensor([[1, 2, 3],
[4, 5, 6]])
torch.Size([2, 3])
tensor([[[1, 2, 3],
[4, 5, 6]]])
torch.Size([1, 2, 3])
tensor([[[1, 2, 3]],
[[4, 5, 6]]])
torch.Size([2, 1, 3])
- 输入四维张量，即input.dim()=4
# 输入：
x = torch.tensor([[[[1, 2, 3],
[4, 5, 6]],
[[0, 2, 1],
[1, 5, 2]]],
[[[1, 2, 3],
[4, 5, 6]],
[[0, 2, 1],
[1, 5, 2]]]])
print(x)
print(x.shape)
y2 = x.unsqueeze(2)
print(y2)
print(y2.shape)
y3 = x.unsqueeze(3)
print(y3)
print(y3.shape)
# 输出：
tensor([[[[1, 2, 3],
[4, 5, 6]],
[[0, 2, 1],
[1, 5, 2]]],
[[[1, 2, 3],
[4, 5, 6]],
[[0, 2, 1],
[1, 5, 2]]]])
torch.Size([2, 2, 2, 3])
tensor([[[[[1, 2, 3],
[4, 5, 6]]],
[[[0, 2, 1],
[1, 5, 2]]]],
[[[[1, 2, 3],
[4, 5, 6]]],
[[[0, 2, 1],
[1, 5, 2]]]]])
torch.Size([2, 2, 1, 2, 3])
tensor([[[[[1, 2, 3]],
[[4, 5, 6]]],
[[[0, 2, 1]],
[[1, 5, 2]]]],
[[[[1, 2, 3]],
[[4, 5, 6]]],
[[[0, 2, 1]],
[[1, 5, 2]]]]])
torch.Size([2, 2, 2, 1, 3])