深入理解Python中的torch.unsqueeze()函数

目录

前言

torch.unsqueeze()函数理解

代码案例

前言

torch.unsqueeze()是PyTorch库中的一个函数，它被用来在输入张量的特定维度上增加一个维度。这是一个很常用的函数，可以用来改变张量的形状和结构。

函数的参数是输入张量和要增加的维度的位置。具体来说，如果你有一个形状为(N, M)的二维张量，torch.unsqueeze(tensor, 0)会在第一维（即N）后面增加一个维度，得到形状为(1, N, M)的张量。如果你在第二维（即M）后面增加一个维度，那么你会得到形状为(N, 1, M)的张量。

这个函数对于很多深度学习任务来说非常重要，比如在处理图像数据时，我们经常需要将二维的图像数据转换为三维的数据，以便于输入到神经网络中。例如，对于一个形状为(N, H, W)的二维图像数据，我们可以通过torch.unsqueeze(tensor, 0)得到形状为(1, N, H, W)的三维数据。

注：

(1, N, H, W) 是三维张量的形状表示，其中 1 表示通道数（Channel），N 表示批量大小（Number），H 表示高度（Height），W 表示宽度（Width）。

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]

代码案例

1. 输入一维张量

# 输入：
x = torch.tensor([1, 2, 3, 4])  
print(x)
print(x.shape)
y = x.unsqueeze(0)
print(y)
print(y.shape)  
z = x.unsqueeze(1)
print(z)
print(z.shape)  

# 输出：
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])

1. 输入二维张量

# 输入：
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])