Python中squeeze函数的用法

squeeze 是 PyTorch 中的一个函数，用于从张量（tensor）中移除所有大小为 1 的维度。这个函数在处理神经网络输出或中间结果时特别有用，因为有时我们可能希望将具有单个元素的维度从张量中移除，以便更容易地进行后续操作。

以下是一些关于如何使用 squeeze 函数的基本示例：

import torch  
  
# 创建一个形状为 (1, 3, 1, 4) 的张量  
x = torch.randn(1, 3, 1, 4)  
print(x.shape)  # 输出: torch.Size([1, 3, 1, 4])  
  
# 使用 squeeze 函数移除所有大小为 1 的维度  
y = x.squeeze()  
print(y.shape)  # 输出: torch.Size([3, 4])  
  
# 也可以指定要移除的维度  
# 例如，只移除第一个维度（索引为 0）  
z = x.squeeze(0)  
print(z.shape)  # 输出: torch.Size([3, 1, 4])  
  
# 或者只移除第三个维度（索引为 2）  
w = x.squeeze(2)  
print(w.shape)  # 输出: torch.Size([1, 3, 4])

注意，squeeze 函数不会改变张量中的元素值，只是改变了张量的形状。如果指定的维度大小不为 1，则 squeeze 函数不会有任何效果。

此外，如果你希望在某些维度上即使它们的大小不为 1 也进行挤压，你可以使用 unsqueeze 函数的逆操作 unsqueeze 来在这些维度上增加大小为 1 的维度，然后再使用 squeeze。但请注意，这通常不是最佳实践，因为它可能会使代码更难理解。在大多数情况下，最好直接处理原始张量的形状，而不是试图通过增加和移除大小为 1 的维度来“修复”它。

作者：焉知有理