pytorch中的squeeze和unsqueeze对比

unsqueeze()

unsqueeze是将Tensor在某一维上增加一个维度

x = torch.randn(2,3)
print(x.shape)
x.unsqueeze_(0)
print(x.shape)
------------------
torch.Size([2, 3])
torch.Size([1, 2, 3])

索引从0开始，也可以在其他维度上增加

x = torch.randn(2,3)
print(x.shape)
x.unsqueeze_(1)
print(x.shape)
--------------------
torch.Size([2, 3])
torch.Size([2, 1, 3])

x = torch.randn(2,3)
print(x.shape)
x.unsqueeze_(2)
print(x.shape)
--------------------
torch.Size([2, 3])
torch.Size([2, 3, 1])

suqeeze()

unsqueeze是将Tensor在某一维度上减少一个维度，且这个维度只能为1

x = torch.randn(2, 1, 3)
print(x.shape)
x.squeeze_(0)
print(x.shape)
x.squeeze_(1)
print(x.shape)
-------------------
torch.Size([2, 1, 3])
torch.Size([2, 1, 3])
torch.Size([2, 3])

任何方法下划线都会改变原始变量，相当于inplace = True，不会创建新的变量，而是直接对原始变量进行修改。