squeeze()的作用

squeeze() 是 Python 的 PyTorch 和 NumPy 等库中常用的一个张量（tensor）操作函数。它的主要作用是去除 tensor 中所有“维度为1”的轴（dimension）。

具体作用

假设有一个张量 x，形状为 (2, 1, 3, 1)，调用 x.squeeze() 后，张量的形状会变成 (2, 3)。也就是说，所有维度为1的轴都被去掉了。

你也可以指定要去掉哪一维（用 squeeze(dim)），比如 x.squeeze(1) 只会去掉第1个轴（从0开始计数），如果该轴的长度正好为1。

示例

import torch

x = torch.randn(2, 1, 3, 1)
print(x.shape)           # 输出: torch.Size([2, 1, 3, 1])

y = x.squeeze()
print(y.shape)           # 输出: torch.Size([2, 3])

z = x.squeeze(1)
print(z.shape)           # 输出: torch.Size([2, 3, 1])

作用总结

• 减少不必要的维度，让张量结构更简洁，方便后续计算。
• 常用于模型输出、损失计算等环节，使数据结构匹配。

如果有更多维度处理的问题，欢迎随时提问！