PyTorch 中的降维操作：torch.squeeze() 函数详解

PyTorch 降维操作：torch.squeeze() 函数详解

在 PyTorch 中，torch.squeeze() 函数可以用于删除指定维度为 1 的维度。下面代码展示了如何利用该函数将输入数据的维度从 4 维降到 3 维：

def forward(self, x):
    x = torch.squeeze(x, dim=3)  # 将输入数据的维度从4维降到3维
    x = self.conv1(x)

代码解析:

• torch.squeeze(x, dim=3)：该语句将 x 的第 4 维（dim=3）删除，即将维度为 (batch_size, channel, height, width) 的输入数据 x 降维为 (batch_size, channel, height)。

总结：

torch.squeeze() 函数是 PyTorch 中常用的降维操作之一，可以方便地移除指定维度为 1 的维度，从而简化数据结构。