torch.squeeze() 这个函数主要对数据的维度进行压缩,去掉维数为1的的维度,默认是将a中所有为1的维度删掉。也可以通过dim指定位置,删掉指定位置的维数为1的维度。
torch.unsqueeze()这个函数主要是对数据维度进行扩充。需要通过dim指定位置,给指定位置加上维数为1的维度。
用法示例:
import torch
x = torch.zeros(3,2,4,1,2,1)# dimension of 3*2*4*1*2
print(x.size()) # torch.Size([3, 2, 4, 1, 2, 1])
print(x.shape)
y = torch.squeeze(x) # Returns a tensor with all the dimensions of input of size 1 removed.
print(y.size()) # torch.Size([3, 2, 4, 2])
print(y.shape)
z = torch.unsqueeze(y,dim=0)# Add a dimension of 1 in the 0th position
print(z.size()) # torch.Size([1, 3, 2, 4, 2])
print(z.shape)
z = torch.unsqueeze(y,dim=1)# Add a dimension of 1 in the 1st position
print(z.size()) # torch.Size([3, 1, 2, 4, 2])
print(z.shape)
y = torch.squeeze(x,dim=0) # remove the 0th position of 1 (no 1)
print('dim=0', y.size()) # torch.Size([3, 2, 4, 1, 2, 1])
print('dim=0', y.shape)
y = torch.squeeze(x, dim=3) # remove the 3rd position of 1 (yes)
print('dim=3', y.size()) # torch.Size([3, 2, 4, 2])
print('dim=3', y.shape)
y = torch.squeeze(x, dim=6) # RuntimeError: Dimension out of range (expected to be in range of [-6, 5], but got 6)
print('dim=6', y.size())
print('dim=6', y.shape)
转载自:
一文掌握torch.squeeze() 和torch.unsqueeze()的用法_荷叶田田-CSDN博客_torch.squeeze