pytorch中squeeze()和unsqueeze()函数的作用是什么

# PyTorch中squeeze()和unsqueeze()函数的作用是什么

在PyTorch中，`squeeze()`和`unsqueeze()`是两个常用的张量维度操作函数，它们主要用于调整张量的维度结构，是神经网络数据处理中的重要工具。

## 1. squeeze()函数

`squeeze()`函数的作用是**移除张量中所有长度为1的维度**（即“压缩”维度），其语法为：

```python
torch.squeeze(input, dim=None) 

• 无参数调用：自动移除所有长度为1的维度

  x = torch.randn(1, 3, 1, 2)
  y = x.squeeze()  # 输出形状变为(3, 2)
  

• 指定dim参数：只移除指定位置的维度（当且仅当该维度长度为1时）

  z = x.squeeze(dim=2)  # 仅移除第2维度，形状变为(1, 3, 2)
  

典型应用场景：当某些网络层（如卷积层）输出包含冗余的单维度时，用于简化张量结构。

2. unsqueeze()函数

unsqueeze()函数的作用是在指定位置插入长度为1的新维度（即”扩展”维度），其语法为：

torch.unsqueeze(input, dim)

示例：

x = torch.tensor([1, 2, 3])
y = x.unsqueeze(0)  # 形状变为(1, 3)
z = x.unsqueeze(1)  # 形状变为(3, 1)

典型应用场景： 1. 为单个数据样本添加batch维度（通常在第0维） 2. 调整张量形状以满足广播运算要求 3. 准备输入特定要求的网络层（如LSTM需要三维输入）

3. 组合使用示例

# 原始张量形状(3,)
x = torch.randn(3)

# 扩展为(1, 3)后矩阵乘法
y = torch.randn(3, 4)
z = x.unsqueeze(0) @ y  # 合法运算

# 运算后移除单维度
result = z.squeeze()

4. 注意事项

1. 对不存在的维度执行操作会报错
2. 在神经网络中，常用于调整数据形状以匹配层输入要求
3. 这两个操作不会改变张量的实际数据，仅修改元数据

理解这两个函数对于正确处理PyTorch中的维度问题至关重要，特别是在数据预处理和网络层间数据传递时经常需要使用。 “`