深入理解PyTorch中的`torch.squeeze()`和`torch.unsqueeze()`函数

在PyTorch中，torch.squeeze()和torch.unsqueeze()是用于处理张量（tensor）的函数，它们允许你添加或删除大小为1的维度。这两个函数在神经网络编程中特别有用，因为它们可以帮助你调整张量的形状以满足模型的需求。

torch.squeeze()

torch.squeeze()函数用于移除张量中所有大小为1的维度。这些维度被称为“轴”（axis）。

语法：

torch.squeeze(input, dim=None)

参数：

• input：输入张量。
• dim：要移除的维度。如果未指定，则默认移除所有大小为1的维度。

示例：
假设我们有一个形状为(2, 1, 3)的张量：

import torch
x = torch.rand(2, 1, 3)
print(x.shape)  # 输出：torch.Size([2, 1, 3])
squeezed = torch.squeeze(x)
print(squeezed.shape)  # 输出：torch.Size([2, 3])

在这个例子中，torch.squeeze()删除了大小为1的中间维度。

torch.unsqueeze()

与torch.squeeze()相反，torch.unsqueeze()函数用于在张量中添加一个大小为1的维度。你可以指定要添加的维度位置。

语法：

torch.unsqueeze(input, dim)

参数：

• input：输入张量。
• dim：要在该位置添加的新维度的维度。

示例：
假设我们有一个形状为(2, 3)的张量：

import torch
x = torch.rand(2, 3)
print(x.shape)  # 输出：torch.Size([2, 3])
unsqueezed = torch.unsqueeze(x, dim=1)
print(unsqueezed.shape)  # 输出：torch.Size([2, 1, 3])

在这个例子中，我们在第二个维度位置（索引为1）添加了一个大小为1的维度。

总结：
torch.squeeze()和torch.unsqueeze()是PyTorch中用于处理张量形状的有用函数。通过合理使用这两个函数，你可以灵活地调整张量的形状以满足不同的模型需求。在实际应用中，这两个函数可以帮助你更好地理解数据在模型中的流动，并确保数据与模型之间的兼容性。