pytorch中的squeeze和unsqueeze的用法小结

1. 理论简介

• squeeze

squeeze的作用是对tensor变量进行维度压缩，去除维数为1的维度。

例如：一个3×2×1×2×1的tensor，squeeze()之后便成了3×2×2。存储的数据并没有发生变化，但是去除了“多余”的维度信息。

• unsqueeze

unsqueeze的作用是对数据维度进行扩充，给指定位置加上维数为1的维度。

例如：原本有个三行的数据，也就是说，它的维度为（3, 1），在0的位置加了一维，那么它的维度变为（1, 3, 1）

请注意，不管是squeeze还是unsqueeze操作，都不会影响原先的张量维度，所以下面的操作中，都将压缩和解压缩操作赋值给了原张量。

2. 举例

2.1 squeeze

（1）x.squeeze()

所有1维数据都进行压缩。

import torch

x = torch.rand(4, 3, 1, 3, 1)
print(x.size())
x = x.squeeze()
print(x.size())

输出：

torch.Size([4, 3, 1, 3, 1])
torch.Size([4, 3, 3])

（2）x.squeeze(n)

指定压缩第n位，如果它的维数为1，则压缩，反之不对该维度操作。

import torch

x = torch.rand(3, 2, 1, 2, 1)
print(x.size()) #torch.Size([3, 2, 1, 2, 1])
x = x.squeeze(0) #第0位的维度为3，不进行压缩
print(x.size()) #torch.Size([3, 2, 1, 2, 1])
x = x.squeeze(2)#第2位的维度为1，进行压缩
print(x.size()) # torch.Size([3, 2, 2, 1])

输出：

torch.Size([3, 2, 1, 2, 1])
torch.Size([3, 2, 1, 2, 1])
torch.Size([3, 2, 2, 1])

（3）torch.squeeze(x)

pytorch中的torch.squeeze(x)方法与上面提到的Tensor对象中的方法作用相同，只不过是语法上有细微的差别。

import torch

x = torch.rand(3, 2, 1, 2, 1)
x = torch.squeeze(x) # 同x.squeeze()
print(x.size())

x = torch.rand(3, 2, 1, 2, 1)
x = torch.squeeze(x, 2) # 同x.squeeze(2),压缩第2维
print(x.size())

输出：

torch.Size([3, 2, 2])
torch.Size([3, 2, 2, 1])

2.2 unsqueeze

（1）x.unsqueeze(n)

x.unsqueeze(n)表示在第n位的位置添加1维。

import torch

x = torch.rand(3, 2)
x = x.unsqueeze(1)
print(x.size())

输出：

torch.Size([3, 1, 2])

（2）torch.unsqueeze(x,n)

torch.unsqueeze(x,n)作用同x.unsqueeze(n)，都表示在第n位的位置添加1维，只不过语法稍有区别。

import torch

x = torch.rand(3, 2)
x = torch.unsqueeze(x, 1) 
print(x.size())

输出：

torch.Size([3, 1, 2])