如何在PyTorch中挤压和展开张量？

PyTorch是一个流行的机器学习库，其中包含许多有用的函数和操作来处理张量数据。有时，我们需要压缩或展开张量以适应特定的数据形状或操作。在本文中，我们将介绍如何使用PyTorch中的squeeze和unsqueeze函数压缩和展开张量。

挤压张量

挤压张量是指从张量中删除大小为1的维度。当我们需要执行某些操作时，这对于高维张量非常有用，这些操作仅需要少量的维度。可以使用PyTorch中的squeeze函数完成此操作。该函数有多种版本，取决于您想要压缩的维度数量。我们将介绍最常用的两个版本。

挤压单一维度：

如果您需要只压缩一个维度，可以使用以下形式的squeeze函数：

import torch

# 创建一个大小为1x3x1x5的张量
x = torch.randn(1, 3, 1, 5)

# 挤压维度2
x_squeezed = torch.squeeze(x, dim=2)

print(x.shape)          # 输出：torch.Size([1, 3, 1, 5])
print(x_squeezed.shape) # 输出：torch.Size([1, 3, 5])

在此示例中，我们创建了一个大小为1x3x1x5的张量，并使用squeeze函数将第二维中的大小1删除。在压缩后，该张量的形状从1x3x1x5变为1x3x5。

挤压所有大小为1的维度：

有时，我们需要删除所有大小为1的维度。这可以通过不提供任何参数调用squeeze函数来完成。以下是一个示例：

import torch

# 创建一个大小为1x3x1x5的张量
x = torch.randn(1, 3, 1, 5)

# 挤压所有大小为1的维度
x_squeezed = torch.squeeze(x)

print(x.shape)          # 输出：torch.Size([1, 3, 1, 5])
print(x_squeezed.shape) # 输出：torch.Size([3, 5])

在上述示例中，我们使用无参的squeeze函数将大小为1的所有维度压缩。

展开张量

展开张量是指将张量的所有维度合并成一个维度。在某些情况下，我们需要将张量展开为一维张量，以便应用某些特定的操作或运算符。可以使用PyTorch中的unsqueeze函数完成此操作。该函数有多种版本，取决于您要展开的维度数量和形状。我们将介绍最常用的两个版本。

展开单一维度：

如果您要展开一个维度，可以使用以下形式的unsqueeze函数：

import torch

# 创建一个大小为1x3x5的张量
x = torch.randn(1, 3, 5)

# 在维度1处展开
x_unsqueezed = torch.unsqueeze(x, dim=1)

print(x.shape)             # 输出：torch.Size([1, 3, 5])
print(x_unsqueezed.shape)  # 输出：torch.Size([1, 1, 3, 5])

在此示例中，我们创建了一个大小为1x3x5的张量，并使用unsqueeze函数将第二维上的大小1展开。展开后，该张量的形状从1x3x5变为1x1x3x5。

展开多个维度：

有时，我们需要展开多个维度，可以同时在多个维度上调用unsqueeze函数。以下是一个示例：

import torch

# 创建一个大小为2x3x2的张量
x = torch.randn(2, 3, 2)

# 在维度0和2处展开
x_unsqueezed = torch.unsqueeze(torch.unsqueeze(x, dim=0), dim=2)

print(x.shape)             # 输出：torch.Size([2, 3, 2])
print(x_unsqueezed.shape)  # 输出：torch.Size([1, 2, 1, 3, 2])

在上述示例中，我们使用了两次unsqueeze函数来在第0和第2维上展开张量。展开后，该张量的形状从2x3x2变为1x2x1x3x2。