pytorch unsqueeze 函数用法
unsqueeze()
是PyTorch中的一个方法,用于在指定位置添加一个额外的维度。这个新添加的维度的大小为1。
函数的定义如下:
torch.Tensor.unsqueeze(dim)
其中,dim
参数是你希望添加新维度的位置。如果原始张量的维度是(A, B, C)
,当你调用.unsqueeze(dim)
时:
- 如果
dim=0
,那么结果张量的大小将会是(1, A, B, C)
- 如果
dim=1
,那么结果张量的大小将会是(A, 1, B, C)
- 如果
dim=2
,那么结果张量的大小将会是(A, B, 1, C)
- 如果
dim=3
,那么结果张量的大小将会是(A, B, C, 1)
注意:dim
可以是负数,表示从后往前数的位置。例如,dim=-1
将在最后一个维度后添加一个维度。
Here is an example:
x = torch.tensor([1, 2, 3, 4]) # Original shape: [4]
y = x.unsqueeze(0) # Shape after unsqueeze: [1, 4]
z = x.unsqueeze(1) # Shape after unsqueeze: [4, 1]
在神经网络应用中,unsqueeze()
方法常被用来满足需要特定输入形状的层或函数的需求,例如需要批处理维度或通道维度的卷积层。