pytorch中的torch.squeeze和torch.unsqueeze

图学习的小张2023-12-28 21:27

python 复制代码

torch.squeeze(input, dim=None) → Tensor

torch.unsqueeze()函数的作用减少数组input指定的维度dim，如果dim不指定，则删除大小为1的维度。返回一个tensor。

如果数组A的维度为（1，1，3）那么执行 torch.squeeze(A，1) 后A的维度变为（1，3）。

python 复制代码

torch.unsqueeze(input, dim) → Tensor

torch.unsqueeze()函数的作用增加数组input指定的维度dim。

例如两行三列的数组A维度为(2，3)，那么这个数组就有三个位置可以增加维度，分别是（ $位置0$ 2， $位置1$ 3 $位置2$ ）或者是（ $位置-3$ 2， $位置-2$ 3 $位置-1$ ），如果执行 torch.unsqueeze(A，1)，数据的维度就变为了（2，1，3）。