pytorch系列代码中常见的两个函数squeeze()和unsqueeze()
1 squeeze
1.1
import torch
x=torch.randn(2,1,2) #生成一个维度是(2,1,2)的tensor向量
print(x.shape) #torch.Size([2, 1, 2])
x=x.squeeze() #参数为空,表示删除tensor元素的所有维度为1的维度,这里的x只有第二个维度是1,所以执行该操作后,第二个维度会被删除剩下:(2,2)
print(x.shape) #torch.Size([2, 2])
x=x.squeeze() #没有维度为1的维度了,所以x的维度保持原样
print(x.shape) #torch.Size([2, 2])
1.2
import torch
x=torch.randn(2,1,2) #生成一个维度是(2,1,2)的tensor向量
print(x.shape) #torch.Size([2, 1, 2])
x=x.squeeze(0) #参数为0,表示如果第0个维度是1维,则删除第0个维度,这里的x只有第1个维度是1,所以执行该操作后,第0个维度不会被删除,x元素的维度依然是(2,1,2)
print(x.shape) #torch.Size([2, 1, 2])
x=x.squeeze(1) #参数为1,表示如果第1个维度是1维,则删除第1个维度,所以执行该操作后,第1个维度会被删除,x元素的维度变成是(2,2)
print(x.shape) #torch.Size([2, 2])
1.3 执行操作后需要写回
import torch
x=torch.randn(2,1,2) #生成一个维度是(2,1,2)的tensor向量
print(x.shape) #torch.Size([2, 1, 2])
x.squeeze() #参数为空,表示删除tensor元素的所有维度为1的维度,这里的x只有第二个维度是1,所以执行该操作后,第二个维度会被删除剩下:(2,2),但是!!!这里对x执行squeeze()操作的结果是返回的,并没有写入原数据空间x,即修改的内容没有保存下来,需要执行x=x.squeeze()才可以
print(x.shape) #torch.Size([2, 1, 2])
2 unsqueeze
import torch
y=torch.randn(2,1,2)
print(y.shape) #torch.Size([2, 1, 2])
y=y.unsqueeze(0) #在第0维,增加一个维度是1的维度
print(y.shape) #torch.Size([1, 2, 1, 2])